An error occurred processing your profile. You can try a lower depth, a larger cutoff, or try profiling a smaller portion of your code. If you continue to have problems you can contact us on GitHub.
| ncalls | tottime | percall | cumtime | percall | filename:lineno(function) |
|---|---|---|---|---|---|
| 47354395 | 1905 | 4.023e-05 | 1971 | 4.162e-05 | array_methods.py:246(deferring_binary_op) |
| 1199000 | 685.1 | 0.0005714 | 690 | 0.0005755 | algorithm.py:337(update_q_profiling) |
| 4189999 | 306.3 | 7.311e-05 | 306.4 | 7.313e-05 | core.py:116(act) |
| 4189999 | 300 | 7.159e-05 | 857.1 | 0.0002046 | algorithm.py:394(get_action) |
| 599500 | 213.6 | 0.0003564 | 215.4 | 0.0003593 | algorithm.py:339(update_pi_profiling) |
| 4200000 | 201.7 | 4.802e-05 | 206.5 | 4.916e-05 | prng.py:1093(threefry_split) |
| 20999995 | 195.1 | 9.293e-06 | 195.1 | 9.293e-06 | ~:0(<method 'step' of 'mujoco_py.cymj.MjSim' objects>) |
| 12599997 | 153 | 1.214e-05 | 157.6 | 1.251e-05 | array.py:351(<genexpr>) |
| 17932402/15852763 | 80.45 | 5.075e-06 | 491.1 | 3.098e-05 | profiler.py:311(wrapper) |
| 2079413 | 74.93 | 3.604e-05 | 74.93 | 3.604e-05 | ~:0(<built-in method jaxlib.xla_extension.execute_sharded>) |
| 14388000 | 74.28 | 5.162e-06 | 1815 | 0.0001261 | update.py:69(<lambda>) |
| 4199999 | 71.81 | 1.71e-05 | 723.8 | 0.0001723 | half_cheetah.py:25(step) |
| 227946981/223746720 | 69.52 | 3.107e-07 | 162.6 | 7.269e-07 | ~:0(<built-in method builtins.isinstance>) |
| 30637520/30637508 | 55.68 | 1.817e-06 | 89.65 | 2.926e-06 | ~:0(<built-in method numpy.core._multiarray_umath.implement_array_function>) |
| 300 | 51.74 | 0.1725 | 1853 | 6.178 | algorithm.py:399(test_agent) |
| 1199000 | 45.34 | 3.782e-05 | 80.11 | 6.681e-05 | algorithm.py:43(sample_batch) |
| 56876624 | 42.96 | 7.553e-07 | 59.7 | 1.05e-06 | config.py:239(get_state) |
| 12599999 | 42.83 | 3.399e-06 | 126.5 | 1.004e-05 | prng.py:100(_check_prng_key_data) |
| 14703901/14703856 | 42.11 | 2.864e-06 | 79.45 | 5.403e-06 | util.py:73(safe_map) |
| 8395998 | 41.73 | 4.97e-06 | 47.43 | 5.649e-06 | array.py:359(is_fully_replicated) |
| 99268417 | 38.38 | 3.866e-07 | 38.38 | 3.866e-07 | array.py:251(shape) |
| 14683169 | 37.19 | 2.533e-06 | 109.7 | 7.475e-06 | core.py:1220(find_top_trace) |
| 98652835 | 35.57 | 3.606e-07 | 61.88 | 6.272e-07 | ~:0(<built-in method _abc._abc_instancecheck>) |
| 4199999 | 34.16 | 8.134e-06 | 266 | 6.332e-05 | mujoco_env.py:224(_step_mujoco_simulation) |
| 134010879 | 32.89 | 2.454e-07 | 58.14 | 4.338e-07 | ~:0(<built-in method builtins.hasattr>) |
| 13771996 | 32.43 | 2.355e-06 | 219.2 | 1.592e-05 | array.py:538(_value) |
| 5420629 | 31.27 | 5.769e-06 | 31.27 | 5.769e-06 | ~:0(<method 'reduce' of 'numpy.ufunc' objects>) |
| 2398000 | 29.2 | 1.218e-05 | 30.4 | 1.268e-05 | core.py:119(set_params) |
| 16194861 | 28.15 | 1.738e-06 | 88.39 | 5.458e-06 | dtypes.py:598(dtype) |
| 4219165 | 24.89 | 5.899e-06 | 256.2 | 6.073e-05 | ~:0(<built-in method numpy.array>) |
| 14683169/14681564 | 24.28 | 1.654e-06 | 771.2 | 5.253e-05 | core.py:362(bind_with_trace) |
| 98652835 | 23.62 | 2.395e-07 | 85.5 | 8.667e-07 | <frozen abc>:117(__instancecheck__) |
| 6274937 | 20.78 | 3.312e-06 | 66.76 | 1.064e-05 | array.py:695(_array_shard_arg) |
| 4189999 | 20.38 | 4.863e-06 | 20.38 | 4.863e-06 | ~:0(<method 'randn' of 'numpy.random.mtrand.RandomState' objects>) |
| 14686370 | 19.86 | 1.352e-06 | 44.39 | 3.023e-06 | ~:0(<built-in method builtins.max>) |
| 4204199 | 19.84 | 4.719e-06 | 54.06 | 1.286e-05 | half_cheetah.py:46(_get_obs) |
| 14682618/14682377 | 19.29 | 1.314e-06 | 911.3 | 6.207e-05 | core.py:357(bind) |
| 32379685/32379635 | 18.8 | 5.808e-07 | 18.84 | 5.818e-07 | ~:0(<built-in method _abc._abc_subclasscheck>) |
| 5998034 | 18.64 | 3.107e-06 | 126 | 2.102e-05 | lax.py:503(_convert_element_type) |
| 34800293 | 18.49 | 5.312e-07 | 39.29 | 1.129e-06 | array.py:259(ndim) |
| 21595857/21594493 | 18.21 | 8.432e-07 | 1834 | 8.494e-05 | tree_util.py:210(<genexpr>) |
| 14683173 | 17.71 | 1.206e-06 | 28.17 | 1.919e-06 | core.py:861(with_cur_sublevel) |
| 3601464/3600391 | 16.51 | 4.585e-06 | 1851 | 0.000514 | ~:0(<built-in method jaxlib.xla_extension.pytree.unflatten>) |
| 18881368 | 16.34 | 8.656e-07 | 35.31 | 1.87e-06 | core.py:450(full_raise) |
| 51581594 | 16.12 | 3.125e-07 | 16.12 | 3.125e-07 | array.py:255(dtype) |
| 1199000 | 15.87 | 1.324e-05 | 2834 | 0.002364 | algorithm.py:334(update) |
| 14680610 | 15.01 | 1.023e-06 | 648.9 | 4.42e-05 | core.py:816(process_primitive) |
| 13759996 | 14.83 | 1.078e-06 | 309.7 | 2.251e-05 | array.py:386(__array__) |
| 5403469 | 14.41 | 2.667e-06 | 14.41 | 2.667e-06 | ~:0(<built-in method jaxlib.xla_extension.pytree.flatten>) |
| 50399996 | 13.88 | 2.754e-07 | 53.6 | 1.064e-06 | prng.py:102(<genexpr>) |
| 24552881 | 13.15 | 5.354e-07 | 18.01 | 7.334e-07 | <frozen importlib._bootstrap>:1209(_handle_fromlist) |
| 1199000 | 12.98 | 1.083e-05 | 34.77 | 2.9e-05 | ~:0(<method 'randint' of 'numpy.random.mtrand.RandomState' objects>) |
| 16216141/16216138 | 12.72 | 7.844e-07 | 69.3 | 4.273e-06 | ~:0(<built-in method builtins.all>) |
| 4199999 | 12.1 | 2.882e-06 | 296.1 | 7.05e-05 | prng.py:602(random_split_impl) |
| 1798523 | 11.79 | 6.553e-06 | 262.4 | 0.0001459 | lax.py:1300(expand_dims) |
| 71489537 | 11.69 | 1.636e-07 | 11.69 | 1.636e-07 | config.py:94(_read) |
| 13762442 | 11.62 | 8.446e-07 | 11.62 | 8.446e-07 | ~:0(<built-in method numpy.asarray>) |
| 4199999 | 11.61 | 2.764e-06 | 534.8 | 0.0001273 | mujoco_env.py:150(do_simulation) |
| 8399998 | 11.52 | 1.371e-06 | 101.8 | 1.212e-05 | prng.py:144(__init__) |
| 4199999 | 11.43 | 2.722e-06 | 742.4 | 0.0001768 | time_limit.py:39(step) |
| 4200001 | 11.43 | 2.721e-06 | 239.7 | 5.707e-05 | random.py:68(_check_prng_key) |
| 6275974/6275303 | 11.42 | 1.82e-06 | 18.56 | 2.957e-06 | functools.py:981(__get__) |
| 1800959/2459 | 11.14 | 0.004531 | 499.4 | 0.2031 | lax_numpy.py:1922(array) |
| 94582573/94582413 | 11.01 | 1.164e-07 | 11.01 | 1.165e-07 | ~:0(<built-in method builtins.len>) |
| 96696981 | 10.34 | 1.069e-07 | 10.34 | 1.069e-07 | ~:0(<method 'get' of 'dict' objects>) |
| 20046933 | 10.34 | 5.159e-07 | 10.34 | 5.159e-07 | ~:0(<built-in method ArrayImpl.>) |
| 4199999 | 10.26 | 2.442e-06 | 51.1 | 1.217e-05 | array.py:345(__iter__) |
| 6275869 | 10.13 | 1.614e-06 | 23.3 | 3.713e-06 | dispatch.py:99(arg_spec) |
| 3600569/3599887 | 9.424 | 2.618e-06 | 1882 | 0.0005228 | tree_util.py:174(tree_map) |
| 2080569 | 9.247 | 4.444e-06 | 283.2 | 0.0001361 | dispatch.py:112(apply_primitive) |
| 2079413 | 9.187 | 4.418e-06 | 192.4 | 9.251e-05 | pxla.py:1901(__call__) |
| 1199999 | 9.116 | 7.596e-06 | 18.38 | 1.531e-05 | algorithm.py:34(store) |
| 20046933 | 9.065 | 4.522e-07 | 19.41 | 9.681e-07 | array.py:508(_check_if_deleted) |
| 6002976 | 8.553 | 1.425e-06 | 21.6 | 3.598e-06 | dtypes.py:375(issubdtype) |
| 14687459 | 8.511 | 5.795e-07 | 24.53 | 1.67e-06 | core.py:1221(<genexpr>) |
| 2080569 | 8.4 | 4.037e-06 | 43.79 | 2.105e-05 | util.py:248(wrapper) |
| 23451803/23449352 | 8.116 | 3.461e-07 | 13.69 | 5.836e-07 | ~:0(<built-in method builtins.getattr>) |
| 16195815 | 7.807 | 4.82e-07 | 28.29 | 1.747e-06 | <__array_function__ internals>:177(result_type) |
| 4200005 | 7.61 | 1.812e-06 | 8.207 | 1.954e-06 | prng.py:119(__instancecheck__) |
| 2081956 | 7.606 | 3.653e-06 | 22.94 | 1.102e-05 | config.py:453(_trace_context) |
| 32379685/32379635 | 7.469 | 2.307e-07 | 26.31 | 8.125e-07 | <frozen abc>:121(__subclasscheck__) |
| 4200000 | 7.153 | 1.703e-06 | 723.2 | 0.0001722 | random.py:205(split) |
| 2079413 | 7.003 | 3.368e-06 | 7.003 | 3.368e-06 | ~:0(<built-in method jaxlib.xla_extension.consume_with_handlers>) |
| 23026705 | 6.754 | 2.933e-07 | 6.754 | 2.933e-07 | array.py:267(sharding) |
| 4200000 | 6.486 | 1.544e-06 | 387.5 | 9.226e-05 | random.py:195(_split) |
| 4200001 | 6.47 | 1.541e-06 | 187.8 | 4.471e-05 | prng.py:733(random_wrap) |
| 11994010 | 5.978 | 4.984e-07 | 27.89 | 2.326e-06 | util.py:325(_arraylike) |
| 13208161 | 5.917 | 4.48e-07 | 17.14 | 1.298e-06 | dtypes.py:124(canonicalize_dtype) |
| 4802843 | 5.895 | 1.227e-06 | 7.493 | 1.56e-06 | sharding_impls.py:343(__eq__) |
| 14683427 | 5.747 | 3.914e-07 | 5.747 | 3.914e-07 | core.py:445(__init__) |
| 8399998 | 5.732 | 6.824e-07 | 15.21 | 1.811e-06 | prng.py:163(shape) |
| 10447181 | 5.67 | 5.427e-07 | 7.958 | 7.617e-07 | config.py:276(get_state) |
| 1798572 | 5.63 | 3.13e-06 | 234.1 | 0.0001302 | lax.py:758(broadcast_in_dim) |
| 4200000 | 5.572 | 1.327e-06 | 88.85 | 2.116e-05 | random.py:83(_return_prng_keys) |
| 4204200 | 5.445 | 1.295e-06 | 34.22 | 8.139e-06 | <__array_function__ internals>:177(concatenate) |
| 1802700 | 5.432 | 3.013e-06 | 6.758 | 3.749e-06 | logx.py:353(store) |
| 6276169 | 5.42 | 8.636e-07 | 77.55 | 1.236e-05 | pxla.py:373(shard_arg) |
| 1798512 | 5.28 | 2.936e-06 | 5.28 | 2.936e-06 | ~:0(<built-in method jaxlib.xla_extension.pytree.flatten_up_to>) |
| 6003063 | 5.201 | 8.664e-07 | 10.97 | 1.828e-06 | numerictypes.py:356(issubdtype) |
| 6280307 | 5.108 | 8.133e-07 | 6.775 | 1.079e-06 | core.py:1523(__hash__) |
| 14683450 | 5.055 | 3.442e-07 | 18.44 | 1.256e-06 | core.py:1214(full_lower) |
| 6276765 | 5.033 | 8.019e-07 | 5.7 | 9.081e-07 | xla.py:194(abstractify) |
| 1798500 | 4.944 | 2.749e-06 | 286.5 | 0.0001593 | lax_numpy.py:866(expand_dims) |
| 1798500 | 4.931 | 2.742e-06 | 1875 | 0.001042 | update.py:47(incremental_update) |
| 28196278 | 4.779 | 1.695e-07 | 4.779 | 1.695e-07 | core.py:1408(is_opaque_dtype) |
| 14683435 | 4.72 | 3.215e-07 | 4.72 | 3.215e-07 | core.py:987(cur_sublevel) |
| 8400004 | 4.719 | 5.618e-07 | 5.668 | 6.748e-07 | prng.py:280(base_arr_shape_to_keys_shape) |
| 4200404 | 4.549 | 1.083e-06 | 10.95 | 2.607e-06 | dispatch.py:257(is_single_device_sharding) |
| 8399998 | 4.527 | 5.389e-07 | 20.34 | 2.421e-06 | prng.py:167(ndim) |
| 1220630 | 4.383 | 3.59e-06 | 17.19 | 1.408e-05 | fromnumeric.py:69(_wrapreduction) |
| 4199999 | 4.256 | 1.013e-06 | 728.1 | 0.0001734 | env_checker.py:33(step) |
| 6017122/6016840 | 4.188 | 6.961e-07 | 22.16 | 3.683e-06 | ~:0(<built-in method builtins.any>) |
| 62 | 4.125 | 0.06652 | 4.125 | 0.06652 | ~:0(<built-in method jaxlib.xla_extension.compile>) |
| 6274926 | 4.043 | 6.444e-07 | 4.838 | 7.71e-07 | sharding_impls.py:340(__hash__) |
| 6276169 | 4.032 | 6.424e-07 | 5.312 | 8.463e-07 | xla.py:160(canonicalize_dtype) |
| 4199999 | 3.845 | 9.155e-07 | 216.5 | 5.155e-05 | prng.py:608(random_split_impl_base) |
| 12006126 | 3.81 | 3.174e-07 | 5.204 | 4.334e-07 | numerictypes.py:282(issubclass_) |
| 2080613 | 3.8 | 1.826e-06 | 81.35 | 3.91e-05 | pxla.py:409(<listcomp>) |
| 4197195 | 3.768 | 8.978e-07 | 29.79 | 7.097e-06 | dtypes.py:457(_dtype_and_weaktype) |
| 3599620 | 3.676 | 1.021e-06 | 8.414 | 2.337e-06 | dtypes.py:669(check_user_dtype_supported) |
| 4199999 | 3.525 | 8.392e-07 | 48.62 | 1.158e-05 | prng.py:745(random_wrap_impl) |
| 4200009 | 3.507 | 8.35e-07 | 7.181 | 1.71e-06 | random.py:103(default_prng_impl) |
| 4200640 | 3.481 | 8.287e-07 | 3.481 | 8.287e-07 | sharding_impls.py:350(device_set) |
| 6274880 | 3.427 | 5.462e-07 | 5.268 | 8.395e-07 | array.py:382(is_fully_addressable) |
| 4200000 | 3.4 | 8.095e-07 | 76.82 | 1.829e-05 | prng.py:762(random_unwrap) |
| 4199999 | 3.315 | 7.892e-07 | 209.8 | 4.995e-05 | prng.py:609(<lambda>) |
| 600 | 3.302 | 0.005504 | 405.7 | 0.6762 | lax_numpy.py:1697(stack) |
| 4199999 | 3.236 | 7.704e-07 | 3.236 | 7.704e-07 | mujoco_env.py:146(dt) |
| 8399998 | 3.153 | 3.753e-07 | 4.567 | 5.437e-07 | array_methods.py:353(_chunk_iter) |
| 10196032 | 3.102 | 3.042e-07 | 3.102 | 3.042e-07 | dtypes.py:581(is_weakly_typed) |
| 5400753 | 3.063 | 5.671e-07 | 17.43 | 3.227e-06 | <__array_function__ internals>:177(ndim) |
| 4200000 | 3.05 | 7.261e-07 | 370.3 | 8.818e-05 | prng.py:591(random_split) |
| 4199999 | 3.033 | 7.22e-07 | 26.76 | 6.371e-06 | ~:0(<method 'sum' of 'numpy.ndarray' objects>) |
| 4200001 | 2.884 | 6.866e-07 | 2.884 | 6.866e-07 | prng.py:513(iterated_vmap_unary) |
| 1800311 | 2.845 | 1.58e-06 | 4.619 | 2.566e-06 | util.py:330(_stackable) |
| 2082456 | 2.843 | 1.365e-06 | 2.843 | 1.365e-06 | ~:0(<built-in method jaxlib.xla_extension.jax_jit.thread_local_state>) |
| 4199999 | 2.842 | 6.767e-07 | 730.9 | 0.000174 | order_enforcing.py:33(step) |
| 9594689 | 2.656 | 2.768e-07 | 16.24 | 1.693e-06 | util.py:339(<genexpr>) |
| 24014913 | 2.594 | 1.08e-07 | 2.594 | 1.08e-07 | ~:0(<built-in method builtins.issubclass>) |
| 2080569 | 2.524 | 1.213e-06 | 204.6 | 9.832e-05 | dispatch.py:211(<lambda>) |
| 3001093 | 2.511 | 8.367e-07 | 2.511 | 8.367e-07 | core.py:1517(__eq__) |
| 3599109 | 2.506 | 6.964e-07 | 11.75 | 3.265e-06 | <__array_function__ internals>:177(shape) |
| 10216075 | 2.469 | 2.417e-07 | 2.469 | 2.417e-07 | ~:0(<built-in method __new__ of type object at 0x860e20>) |
| 1823 | 2.385 | 0.001308 | 34.49 | 0.01892 | dtypes.py:619(_lattice_result_type) |
| 2079413 | 2.374 | 1.142e-06 | 8.191 | 3.939e-06 | dispatch.py:430(needs_check_special) |
| 5400753 | 2.326 | 4.306e-07 | 8.531 | 1.58e-06 | fromnumeric.py:3152(ndim) |
| 1798556 | 2.255 | 1.254e-06 | 88.58 | 4.925e-05 | lax_numpy.py:2008(asarray) |
| 19773391 | 2.251 | 1.139e-07 | 2.251 | 1.139e-07 | typing.py:2233(cast) |
| 2079413 | 2.051 | 9.866e-07 | 2.051 | 9.866e-07 | pxla.py:1903(<listcomp>) |
| 18876311 | 1.941 | 1.028e-07 | 1.941 | 1.028e-07 | core.py:813(pure) |
| 3600569 | 1.931 | 5.362e-07 | 7.211 | 2.003e-06 | tree_util.py:209(<listcomp>) |
| 2080613 | 1.896 | 9.115e-07 | 83.24 | 4.001e-05 | pxla.py:388(shard_args) |
| 4165538 | 1.864 | 4.474e-07 | 2.62 | 6.29e-07 | config.py:445(get_state) |
| 12568037/12567759 | 1.859 | 1.479e-07 | 1.859 | 1.479e-07 | ~:0(<built-in method builtins.hash>) |
| 1798627 | 1.839 | 1.022e-06 | 5.095 | 2.833e-06 | api_util.py:68(_ensure_index_tuple) |
| 4199999 | 1.813 | 4.316e-07 | 23.73 | 5.649e-06 | _methods.py:47(_sum) |
| 1199008 | 1.724 | 1.438e-06 | 18.79 | 1.567e-05 | fromnumeric.py:2955(prod) |
| 2080613 | 1.694 | 8.141e-07 | 91.07 | 4.377e-05 | pxla.py:1674(__call__) |
| 4199999 | 1.692 | 4.029e-07 | 2.446 | 5.825e-07 | prng.py:775(random_unwrap_impl) |
| 16788780 | 1.692 | 1.008e-07 | 1.692 | 1.008e-07 | dtypes.py:459(<genexpr>) |
| 3599109 | 1.607 | 4.465e-07 | 3.297 | 9.16e-07 | fromnumeric.py:1991(shape) |
| 1199008 | 1.539 | 1.283e-06 | 21.78 | 1.817e-05 | <__array_function__ internals>:177(prod) |
| 3602508 | 1.502 | 4.169e-07 | 14.53 | 4.034e-06 | tree_util.py:42(tree_flatten) |
| 8399998 | 1.482 | 1.765e-07 | 1.482 | 1.765e-07 | prng.py:151(unsafe_raw_array) |
| 1800680 | 1.474 | 8.184e-07 | 19.65 | 1.091e-05 | util.py:336(check_arraylike) |
| 6002976 | 1.443 | 2.404e-07 | 2.077 | 3.46e-07 | dtypes.py:364(_issubclass) |
| 1805914/1805912 | 1.412 | 7.821e-07 | 3.017 | 1.671e-06 | core.py:1372(concrete_or_error) |
| 16195815 | 1.389 | 8.574e-08 | 1.389 | 8.574e-08 | multiarray.py:668(result_type) |
| 1800329 | 1.319 | 7.329e-07 | 1.496 | 8.312e-07 | util.py:327(canonicalize_axis) |
| 4199018 | 1.305 | 3.107e-07 | 31.09 | 7.404e-06 | dtypes.py:620(<genexpr>) |
| 1798523 | 1.301 | 7.233e-07 | 2.794 | 1.554e-06 | lax.py:1305(<listcomp>) |
| 1255 | 1.294 | 0.001031 | 85.33 | 0.068 | util.py:284(<listcomp>) |
| 1800053 | 1.277 | 7.095e-07 | 1.277 | 7.095e-07 | ~:0(<built-in method builtins.sorted>) |
| 9322987 | 1.177 | 1.263e-07 | 1.177 | 1.263e-07 | ~:0(<method 'items' of 'dict' objects>) |
| 3600622 | 1.1 | 3.054e-07 | 1.1 | 3.054e-07 | util.py:331(<genexpr>) |
| 1220630 | 1.082 | 8.863e-07 | 1.082 | 8.863e-07 | fromnumeric.py:70(<dictcomp>) |
| 9702731 | 1.066 | 1.099e-07 | 1.066 | 1.099e-07 | ~:0(<built-in method builtins.id>) |
| 4210003 | 1.033 | 2.454e-07 | 1.033 | 2.454e-07 | box.py:136(shape) |
| 1800961 | 0.9588 | 5.324e-07 | 2.341 | 1.3e-06 | tree_util.py:80(tree_leaves) |
| 6275974 | 0.9255 | 1.475e-07 | 0.9255 | 1.475e-07 | ~:0(<method '__exit__' of '_thread.RLock' objects>) |
| 6039962 | 0.8979 | 1.487e-07 | 0.8979 | 1.487e-07 | ~:0(<method 'append' of 'list' objects>) |
| 3606059 | 0.8375 | 2.323e-07 | 1.203 | 3.336e-07 | lax_numpy.py:1945(<lambda>) |
| 3609118 | 0.8157 | 2.26e-07 | 2.477 | 6.864e-07 | lax_numpy.py:1968(<genexpr>) |
| 600 | 0.8047 | 0.001341 | 89.38 | 0.149 | lax_numpy.py:1980(<listcomp>) |
| 4204200 | 0.7696 | 1.831e-07 | 0.7696 | 1.831e-07 | multiarray.py:152(concatenate) |
| 1200047 | 0.765 | 6.375e-07 | 0.765 | 6.375e-07 | ~:0(<built-in method builtins.min>) |
| 1270 | 0.7489 | 0.0005897 | 0.9864 | 0.0007767 | dtypes.py:635(<setcomp>) |
| 6278988 | 0.7337 | 1.169e-07 | 0.7337 | 1.169e-07 | ~:0(<method 'values' of 'dict' objects>) |
| 1207 | 0.717 | 0.000594 | 242.5 | 0.2009 | lax_numpy.py:1750(concatenate) |
| 280276 | 0.684 | 2.44e-06 | 113.9 | 0.0004065 | lax.py:580(concatenate) |
| 5400753 | 0.6817 | 1.262e-07 | 0.6817 | 1.262e-07 | fromnumeric.py:3148(_ndim_dispatcher) |
| 1798523 | 0.6757 | 3.757e-07 | 0.6757 | 3.757e-07 | lax.py:1312(<listcomp>) |
| 6274937 | 0.6342 | 1.011e-07 | 0.6342 | 1.011e-07 | pxla.py:371(identity) |
| 3599109 | 0.5563 | 1.546e-07 | 0.5563 | 1.546e-07 | fromnumeric.py:1987(_shape_dispatcher) |
| 2411 | 0.5028 | 0.0002085 | 1.049 | 0.0004349 | reductions.py:319(mean) |
| 3607 | 0.4912 | 0.0001362 | 114.4 | 0.03172 | lax_numpy.py:1774(<listcomp>) |
| 4205024 | 0.4587 | 1.091e-07 | 0.4587 | 1.091e-07 | ~:0(<method 'keys' of 'dict' objects>) |
| 3614795 | 0.4342 | 1.201e-07 | 0.4342 | 1.201e-07 | ~:0(<built-in method _operator.index>) |
| 1798513 | 0.4186 | 2.328e-07 | 0.4186 | 2.328e-07 | ~:0(<method 'insert' of 'list' objects>) |
| 4200 | 0.4079 | 9.712e-05 | 0.4079 | 9.712e-05 | ~:0(<method 'reset' of 'mujoco_py.cymj.MjSim' objects>) |
| 1199008 | 0.2633 | 2.196e-07 | 0.2633 | 2.196e-07 | fromnumeric.py:2950(_prod_dispatcher) |
| 2400 | 0.2555 | 0.0001065 | 1.17 | 0.0004876 | reductions.py:480(std) |
| 4196644 | 0.2375 | 5.658e-08 | 0.2375 | 5.659e-08 | dtypes.py:447(_jax_type) |
| 10000 | 0.2262 | 2.262e-05 | 0.6921 | 6.921e-05 | box.py:171(sample) |
| 14200 | 0.2103 | 1.481e-05 | 0.4071 | 2.867e-05 | ~:0(<method 'uniform' of 'numpy.random._generator.Generator' objects>) |
| 1798662 | 0.2043 | 1.136e-07 | 0.2043 | 1.136e-07 | core.py:1375(<lambda>) |
| 2400 | 0.1008 | 4.199e-05 | 502.3 | 0.2093 | logx.py:28(mpi_statistics_scalar) |
| 1256 | 0.06465 | 5.148e-05 | 119.9 | 0.09545 | util.py:276(promote_dtypes) |
| 4200 | 0.04864 | 1.158e-05 | 0.04864 | 1.158e-05 | ~:0(<method 'forward' of 'mujoco_py.cymj.MjSim' objects>) |
| 4200 | 0.04571 | 1.088e-05 | 0.3736 | 8.896e-05 | half_cheetah.py:54(reset_model) |
| 21300/300 | 0.04445 | 0.0001482 | 0.1688 | 0.0005627 | pickle.py:535(save) |
| 53428 | 0.04305 | 8.058e-07 | 0.07793 | 1.459e-06 | __init__.py:97(_xla_gc_callback) |
| 10000 | 0.0427 | 4.27e-06 | 0.0427 | 4.27e-06 | ~:0(<method 'normal' of 'numpy.random._generator.Generator' objects>) |
| 1200 | 0.04158 | 3.465e-05 | 0.2923 | 0.0002436 | reductions.py:271(min) |
| 300 | 0.03922 | 0.0001307 | 0.7855 | 0.002618 | logx.py:295(dump_tabular) |
| 46 | 0.03855 | 0.0008381 | 0.03876 | 0.0008427 | lax.py:3070(_compute_squeeze_shape) |
| 1200 | 0.03842 | 3.202e-05 | 0.289 | 0.0002409 | reductions.py:256(max) |
| 7500 | 0.03814 | 5.085e-06 | 0.03814 | 5.085e-06 | ~:0(<built-in method builtins.print>) |
| 53428 | 0.03488 | 6.528e-07 | 0.03488 | 6.528e-07 | ~:0(<built-in method jaxlib.xla_extension.collect_garbage>) |
| 3300 | 0.02732 | 8.279e-06 | 640 | 0.1939 | logx.py:365(log_tabular) |
| 20000 | 0.02533 | 1.267e-06 | 0.02533 | 1.267e-06 | ~:0(<method 'exponential' of 'numpy.random._generator.Generator' objects>) |
| 4200 | 0.02389 | 5.689e-06 | 0.09092 | 2.165e-05 | ~:0(<method 'get_state' of 'mujoco_py.cymj.MjSim' objects>) |
| 8400 | 0.02214 | 2.636e-06 | 0.05248 | 6.248e-06 | copy.py:128(deepcopy) |
| 139/62 | 0.02122 | 0.0003422 | 0.2692 | 0.004342 | mlir.py:1098(jaxpr_subcomp) |
| 4200 | 0.02023 | 4.816e-06 | 0.2211 | 5.265e-05 | mujoco_env.py:217(set_state) |
| 1736 | 0.01945 | 1.12e-05 | 0.01945 | 1.12e-05 | ~:0(<built-in method >) |
| 4200 | 0.01897 | 4.517e-06 | 0.05427 | 1.292e-05 | ~:0(<method 'set_state' of 'mujoco_py.cymj.MjSim' objects>) |
| 8543 | 0.01829 | 2.14e-06 | 0.01969 | 2.305e-06 | util.py:51(safe_zip) |
| 4200 | 0.01641 | 3.907e-06 | 0.8022 | 0.000191 | mujoco_env.py:125(reset) |
| 4200 | 0.01614 | 3.843e-06 | 0.01614 | 3.843e-06 | ~:0(<method 'standard_normal' of 'numpy.random._generator.Generator' objects>) |
| 30000/10000 | 0.01604 | 1.604e-06 | 0.01604 | 1.604e-06 | core.py:253(action_space) |
| 300 | 0.01598 | 5.327e-05 | 0.01598 | 5.327e-05 | ~:0(<built-in method io.open>) |
| 10800 | 0.01513 | 1.401e-06 | 0.02348 | 2.174e-06 | pickle.py:491(memoize) |
| 21300/300 | 0.01472 | 4.905e-05 | 0.1696 | 0.0005654 | numpy_pickle.py:322(save) |
| 11348 | 0.01446 | 1.274e-06 | 0.09319 | 8.212e-06 | fromnumeric.py:2432(all) |
| 11348 | 0.01441 | 1.27e-06 | 0.1206 | 1.063e-05 | <__array_function__ internals>:177(all) |
| 10057 | 0.01438 | 1.43e-06 | 0.01438 | 1.43e-06 | ~:0(<built-in method numpy.empty>) |
| 300 | 0.01419 | 4.731e-05 | 0.01419 | 4.731e-05 | ~:0(<method '__exit__' of '_io._IOBase' objects>) |
| 9520 | 0.01335 | 1.403e-06 | 0.02617 | 2.749e-06 | core.py:701(__getattr__) |
| 8400 | 0.01314 | 1.564e-06 | 0.01457 | 1.735e-06 | copy.py:243(_keep_alive) |
| 545/46 | 0.01293 | 0.0002811 | 0.7202 | 0.01566 | pjit.py:410(common_infer_params) |
| 16800 | 0.01278 | 7.607e-07 | 0.04413 | 2.627e-06 | <__array_function__ internals>:177(copy) |
| 1881 | 0.01087 | 5.781e-06 | 0.207 | 0.00011 | partial_eval.py:1793(default_process_primitive) |
| 1756 | 0.01074 | 6.114e-06 | 0.1494 | 8.508e-05 | utils.py:56(standard_abstract_eval) |
| 6600 | 0.01073 | 1.626e-06 | 0.03136 | 4.751e-06 | pickle.py:851(save_str) |
| 10035 | 0.01039 | 1.035e-06 | 0.02918 | 2.908e-06 | core.py:1492(__init__) |
| 36600 | 0.01033 | 2.821e-07 | 0.0145 | 3.962e-07 | pickle.py:241(write) |
| 6900 | 0.01025 | 1.486e-06 | 0.01025 | 1.486e-06 | logx.py:130(log_tabular) |
| 6000 | 0.0101 | 1.684e-06 | 0.6335 | 0.0001056 | array.py:290(__float__) |
| 10012 | 0.00983 | 9.818e-07 | 0.00983 | 9.818e-07 | ~:0(<method 'astype' of 'numpy.ndarray' objects>) |
| 62 | 0.00976 | 0.0001574 | 0.00976 | 0.0001574 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlirHlo.register_mhlo_dialect>) |
| 10468 | 0.0097 | 9.266e-07 | 0.01437 | 1.372e-06 | core.py:2026(canonicalize_shape) |
| 6000 | 0.009581 | 1.597e-06 | 0.01959 | 3.265e-06 | arrayprint.py:506(wrapper) |
| 492 | 0.009526 | 1.936e-05 | 0.2577 | 0.0005237 | core.py:410(eval_jaxpr) |
| 10274 | 0.009337 | 9.088e-07 | 0.08746 | 8.513e-06 | <__array_function__ internals>:177(any) |
| 10274 | 0.00932 | 9.071e-07 | 0.05249 | 5.109e-06 | fromnumeric.py:2333(any) |
| 21600 | 0.009291 | 4.302e-07 | 0.01206 | 5.582e-07 | pickle.py:217(commit_frame) |
| 1200 | 0.009134 | 7.612e-06 | 0.0959 | 7.992e-05 | pxla.py:3254(_execute_trivial) |
| 1050 | 0.009131 | 8.696e-06 | 0.01216 | 1.159e-05 | dataclasses.py:1446(replace) |
| 10911 | 0.008927 | 8.182e-07 | 0.02237 | 2.05e-06 | __init__.py:450(_replace) |
| 1800/300 | 0.008782 | 2.927e-05 | 0.164 | 0.0005466 | pickle.py:978(_batch_setitems) |
| 62 | 0.008573 | 0.0001383 | 0.008573 | 0.0001383 | ~:0(<built-in method jaxlib.mlir._mlir_libs._stablehlo.register_dialect>) |
| 8400 | 0.008263 | 9.837e-07 | 0.01028 | 1.224e-06 | copy.py:227(_deepcopy_dict) |
| 971 | 0.008014 | 8.253e-06 | 0.04738 | 4.88e-05 | mlir.py:322(_source_info_to_location) |
| 801 | 0.007971 | 9.952e-06 | 0.02284 | 2.851e-05 | inspect.py:2331(_signature_from_function) |
| 62 | 0.007858 | 0.0001267 | 0.01228 | 0.0001981 | __init__.py:94(__init__) |
| 7244 | 0.007539 | 1.041e-06 | 0.02476 | 3.418e-06 | reductions.py:187(_ensure_optional_axes) |
| 6000 | 0.007331 | 1.222e-06 | 0.02692 | 4.487e-06 | arrayprint.py:1571(_array_str_implementation) |
| 2372 | 0.007172 | 3.024e-06 | 0.01747 | 7.364e-06 | core.py:2519(subst_axis_names_var) |
| 40000 | 0.007171 | 1.793e-07 | 0.007171 | 1.793e-07 | space.py:72(np_random) |
| 4200 | 0.006879 | 1.638e-06 | 0.8093 | 0.0001927 | env_checker.py:41(reset) |
| 20823/19388 | 0.006841 | 3.529e-07 | 0.0573 | 2.955e-06 | ~:0(<built-in method builtins.next>) |
| 6000 | 0.006796 | 1.133e-06 | 0.05575 | 9.291e-06 | array.py:275(__str__) |
| 77/62 | 0.006758 | 0.000109 | 0.3007 | 0.004849 | mlir.py:850(lower_jaxpr_to_fun) |
| 16800 | 0.006653 | 3.96e-07 | 0.01992 | 1.186e-06 | function_base.py:871(copy) |
| 8299 | 0.006108 | 7.36e-07 | 0.01396 | 1.682e-06 | source_info_util.py:126(is_user_filename) |
| 1200 | 0.005814 | 4.845e-06 | 0.03541 | 2.951e-05 | pickle.py:1056(save_global) |
| 4286/4232 | 0.005735 | 1.355e-06 | 0.06686 | 1.58e-05 | ~:0(<method 'join' of 'str' objects>) |
| 10035 | 0.005721 | 5.701e-07 | 0.005721 | 5.701e-07 | core.py:1418(_dtype_object) |
| 1942 | 0.005623 | 2.895e-06 | 0.01353 | 6.966e-06 | partial_eval.py:1774(makevar) |
| 250 | 0.005588 | 2.235e-05 | 0.02918 | 0.0001167 | partial_eval.py:1674(_inline_literals) |
| 250 | 0.005567 | 2.227e-05 | 0.01644 | 6.576e-05 | partial_eval.py:1630(_const_folding_and_forwarding) |
| 23 | 0.005542 | 0.000241 | 1.064 | 0.04627 | dispatch.py:199(xla_primitive_callable) |
| 300 | 0.005537 | 1.846e-05 | 0.2162 | 0.0007208 | numpy_pickle.py:424(dump) |
| 4200 | 0.005475 | 1.304e-06 | 0.8194 | 0.0001951 | time_limit.py:58(reset) |
| 1008 | 0.005423 | 5.38e-06 | 0.04039 | 4.007e-05 | lax.py:1514(naryop_dtype_rule) |
| 1200/300 | 0.005398 | 1.799e-05 | 0.1544 | 0.0005147 | pickle.py:621(save_reduce) |
| 4240 | 0.005395 | 1.272e-06 | 0.006197 | 1.462e-06 | contextlib.py:104(__init__) |
| 300 | 0.005317 | 1.772e-05 | 0.005317 | 1.772e-05 | ~:0(<method 'flush' of '_io.TextIOWrapper' objects>) |
| 6000 | 0.005308 | 8.846e-07 | 0.00645 | 1.075e-06 | arrayprint.py:1564(_guarded_repr_or_str) |
| 10911 | 0.005218 | 4.782e-07 | 0.01343 | 1.231e-06 | __init__.py:440(_make) |
| 299 | 0.005072 | 1.696e-05 | 0.01145 | 3.83e-05 | _stablehlo_ops_gen.py:1199(__init__) |
| 4200 | 0.005053 | 1.203e-06 | 0.005053 | 1.203e-06 | mujoco_env.py:140(set_state) |
| 802/801 | 0.004836 | 6.037e-06 | 0.02979 | 3.72e-05 | inspect.py:2426(_signature_from_callable) |
| 998 | 0.004741 | 4.75e-06 | 0.004741 | 4.75e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.inferReturnTypes>) |
| 64 | 0.004685 | 7.32e-05 | 0.4066 | 0.006354 | pxla.py:2371(lower_sharding_computation) |
| 4200 | 0.004647 | 1.106e-06 | 0.8139 | 0.0001938 | order_enforcing.py:39(reset) |
| 1232 | 0.004637 | 3.764e-06 | 0.03581 | 2.906e-05 | pxla.py:446(batched_device_put) |
| 2676 | 0.004589 | 1.715e-06 | 0.007819 | 2.922e-06 | inspect.py:2683(__init__) |
| 1232 | 0.004535 | 3.681e-06 | 0.05146 | 4.177e-05 | pxla.py:426(_shard_array) |
| 545 | 0.004462 | 8.188e-06 | 0.006112 | 1.121e-05 | inspect.py:3075(_bind) |
| 4192/3654 | 0.004325 | 1.184e-06 | 0.01679 | 4.595e-06 | contextlib.py:141(__exit__) |
| 6155 | 0.00432 | 7.018e-07 | 0.009981 | 1.622e-06 | core.py:300(<lambda>) |
| 7209 | 0.004241 | 5.882e-07 | 0.005429 | 7.53e-07 | lax_numpy.py:145(__hash__) |
| 62 | 0.004224 | 6.813e-05 | 0.004224 | 6.813e-05 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.load_all_available_dialects>) |
| 6619/3580 | 0.004187 | 1.17e-06 | 0.007911 | 2.21e-06 | tree_util.py:706(_generate_key_paths_) |
| 8529 | 0.004172 | 4.891e-07 | 0.006283 | 7.367e-07 | core.py:259(__init__) |
| 566/551 | 0.004166 | 7.56e-06 | 0.0426 | 7.732e-05 | core.py:2565(subst_axis_names_jaxpr) |
| 4770 | 0.004119 | 8.636e-07 | 0.01139 | 2.389e-06 | source_info_util.py:205(user_context) |
| 62 | 0.004071 | 6.566e-05 | 0.004071 | 6.566e-05 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.verify>) |
| 36300 | 0.004001 | 1.102e-07 | 0.004001 | 1.102e-07 | ~:0(<method 'write' of '_io.BytesIO' objects>) |
| 5174 | 0.003948 | 7.63e-07 | 0.01651 | 3.19e-06 | source_info_util.py:117(replace) |
| 545 | 0.003944 | 7.236e-06 | 0.04468 | 8.199e-05 | api_util.py:611(_arg_names) |
| 62 | 0.003926 | 6.333e-05 | 0.003926 | 6.333e-05 | ~:0(<built-in method jaxlib.mlir._mlir_libs._chlo.register_dialect>) |
| 1200 | 0.003586 | 2.989e-06 | 0.003586 | 2.989e-06 | pxla.py:1696(<listcomp>) |
| 216 | 0.003572 | 1.653e-05 | 0.008308 | 3.846e-05 | _stablehlo_ops_gen.py:61(__init__) |
| 10990 | 0.003559 | 3.238e-07 | 0.004666 | 4.245e-07 | core.py:2010(_canonicalize_dimension) |
| 504 | 0.003522 | 6.987e-06 | 0.008612 | 1.709e-05 | lax.py:2373(<genexpr>) |
| 21423 | 0.003521 | 1.644e-07 | 0.003521 | 1.644e-07 | partial_eval.py:1719(frame) |
| 2369/1527 | 0.003414 | 2.236e-06 | 0.2639 | 0.0001728 | partial_eval.py:1788(process_primitive) |
| 200 | 0.003392 | 1.696e-05 | 0.007709 | 3.854e-05 | _stablehlo_ops_gen.py:3037(__init__) |
| 4200 | 0.003379 | 8.046e-07 | 0.4113 | 9.793e-05 | mujoco_env.py:214(_reset_simulation) |
| 250/113 | 0.003365 | 2.978e-05 | 0.7226 | 0.006394 | partial_eval.py:2054(trace_to_subjaxpr_dynamic) |
| 274 | 0.003331 | 1.216e-05 | 0.02911 | 0.0001062 | mlir.py:257(_ndarray_constant_handler) |
| 971 | 0.003306 | 3.405e-06 | 0.003306 | 3.405e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.file>) |
| 971 | 0.003303 | 3.401e-06 | 0.003303 | 3.401e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.name>) |
| 62 | 0.003277 | 5.286e-05 | 0.3542 | 0.005712 | mlir.py:657(lower_jaxpr_to_module) |
| 1444 | 0.003272 | 2.266e-06 | 0.01211 | 8.388e-06 | xla.py:216(_make_shaped_array_for_numpy_array) |
| 20083 | 0.003265 | 1.626e-07 | 0.004748 | 2.364e-07 | source_info_util.py:129(<genexpr>) |
| 1800/300 | 0.003247 | 1.082e-05 | 0.1664 | 0.0005546 | pickle.py:965(save_dict) |
| 1404 | 0.003237 | 2.305e-06 | 0.003237 | 2.305e-06 | ~:0(<built-in method jaxlib.xla_extension.get_traceback>) |
| 545/44 | 0.003137 | 7.13e-05 | 4.272 | 0.09709 | pjit.py:179(_python_pjit_helper) |
| 1032 | 0.003127 | 3.03e-06 | 0.01236 | 1.198e-05 | partial_eval.py:1723(new_arg) |
| 547/46 | 0.003119 | 6.781e-05 | 0.6996 | 0.01521 | linear_util.py:308(memoized_fun) |
| 62 | 0.003113 | 5.021e-05 | 0.003113 | 5.021e-05 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.write_bytecode>) |
| 2131/551 | 0.003108 | 5.641e-06 | 0.04612 | 8.371e-05 | core.py:2495(subst_axis_names) |
| 840 | 0.003107 | 3.699e-06 | 0.009891 | 1.177e-05 | lax.py:1537(broadcasting_shape_rule) |
| 5147 | 0.002988 | 5.806e-07 | 0.01106 | 2.149e-06 | core.py:242(replace) |
| 11100 | 0.002936 | 2.645e-07 | 0.002936 | 2.645e-07 | ~:0(<built-in method _struct.pack>) |
| 1756 | 0.0029 | 1.651e-06 | 0.005624 | 3.203e-06 | utils.py:114(standard_named_shape_rule) |
| 3211 | 0.002871 | 8.942e-07 | 0.004616 | 1.438e-06 | partial_eval.py:1768(getvar) |
| 9245 | 0.002809 | 3.039e-07 | 0.004056 | 4.387e-07 | core.py:1775(raise_to_shaped) |
| 1200 | 0.002729 | 2.274e-06 | 0.003115 | 2.596e-06 | ~:0(<method '__reduce_ex__' of 'object' objects>) |
| 1397 | 0.002729 | 1.953e-06 | 0.01316 | 9.421e-06 | mlir.py:142(_array_ir_types) |
| 1580 | 0.002684 | 1.699e-06 | 0.02264 | 1.433e-05 | core.py:2534(subst_axis_names_eqn) |
| 8400 | 0.002675 | 3.184e-07 | 0.002777 | 3.306e-07 | core.py:75(np_random) |
| 4192 | 0.002653 | 6.328e-07 | 0.008815 | 2.103e-06 | contextlib.py:287(helper) |
| 1736 | 0.002638 | 1.52e-06 | 0.002638 | 1.52e-06 | _ods_common.py:114(get_default_loc_context) |
| 581 | 0.002604 | 4.482e-06 | 0.04492 | 7.732e-05 | lax.py:1653(_nary_lower_hlo) |
| 3217 | 0.002566 | 7.976e-07 | 0.009988 | 3.105e-06 | source_info_util.py:185(current) |
| 11348 | 0.002565 | 2.261e-07 | 0.002565 | 2.261e-07 | fromnumeric.py:2427(_all_dispatcher) |
| 62 | 0.002563 | 4.133e-05 | 4.147 | 0.06689 | pxla.py:2950(from_hlo) |
| 545/46 | 0.002551 | 5.545e-05 | 0.7206 | 0.01567 | api.py:304(infer_params) |
| 802 | 0.002496 | 3.112e-06 | 0.004024 | 5.018e-06 | inspect.py:2969(__init__) |
| 584 | 0.002495 | 4.272e-06 | 0.01995 | 3.416e-05 | mlir.py:1310(multi_broadcast_in_dim) |
| 77/64 | 0.002492 | 3.894e-05 | 0.009908 | 0.0001548 | partial_eval.py:1357(_dce_jaxpr) |
| 5823 | 0.002481 | 4.261e-07 | 0.009666 | 1.66e-06 | partial_eval.py:1688(<lambda>) |
| 600 | 0.00247 | 4.116e-06 | 0.1083 | 0.0001805 | lax.py:4289(_copy_impl) |
| 254/113 | 0.002469 | 2.185e-05 | 0.6867 | 0.006077 | linear_util.py:152(call_wrapped) |
| 500 | 0.002463 | 4.927e-06 | 0.005506 | 1.101e-05 | config.py:606(update_thread_local_jit_state) |
| 2530 | 0.002442 | 9.653e-07 | 0.003993 | 1.578e-06 | _ods_common.py:126(get_op_result_or_value) |
| 2010 | 0.002435 | 1.212e-06 | 0.006051 | 3.01e-06 | core.py:246(new_jaxpr_eqn) |
| 1255 | 0.002423 | 1.931e-06 | 0.008953 | 7.134e-06 | api_util.py:569(_numpy_array_abstractify) |
| 250 | 0.002395 | 9.58e-06 | 0.04928 | 0.0001971 | partial_eval.py:1579(to_jaxpr) |
| 610 | 0.002391 | 3.92e-06 | 0.01787 | 2.929e-05 | source_info_util.py:157(<genexpr>) |
| 1894 | 0.002387 | 1.26e-06 | 0.002387 | 1.26e-06 | util.py:418(__eq__) |
| 1978 | 0.002347 | 1.186e-06 | 0.1635 | 8.267e-05 | core.py:395(abstract_eval_) |
| 3164 | 0.002331 | 7.367e-07 | 0.003162 | 9.994e-07 | partial_eval.py:1476(__init__) |
| 1696 | 0.002289 | 1.349e-06 | 0.005194 | 3.063e-06 | mlir.py:132(dtype_to_ir_type) |
| 4192/3654 | 0.002271 | 6.216e-07 | 0.02426 | 6.638e-06 | contextlib.py:132(__enter__) |
| 504 | 0.002231 | 4.427e-06 | 0.005471 | 1.086e-05 | lax.py:2380(<genexpr>) |
| 16800 | 0.002198 | 1.308e-07 | 0.002198 | 1.308e-07 | function_base.py:867(_copy_dispatcher) |
| 1058 | 0.002159 | 2.04e-06 | 0.003625 | 3.426e-06 | source_info_util.py:99(__str__) |
| 62 | 0.002147 | 3.464e-05 | 0.005253 | 8.473e-05 | xla_bridge.py:89(get_compile_options) |
| 3754 | 0.002142 | 5.705e-07 | 0.01037 | 2.762e-06 | core.py:1498(update) |
| 299 | 0.002141 | 7.162e-06 | 0.01763 | 5.896e-05 | mlir.py:231(_numpy_array_constant) |
| 3164 | 0.002127 | 6.723e-07 | 0.008889 | 2.81e-06 | partial_eval.py:1607(newvar) |
| 6076 | 0.002124 | 3.496e-07 | 0.01832 | 3.015e-06 | lax.py:1517(<genexpr>) |
| 545/44 | 0.002112 | 4.8e-05 | 4.273 | 0.09712 | pjit.py:236(cache_miss) |
| 2400 | 0.00211 | 8.791e-07 | 0.004929 | 2.054e-06 | array.py:278(__len__) |
| 1051 | 0.002078 | 1.977e-06 | 0.01207 | 1.148e-05 | lax.py:4535(_check_same_dtypes) |
| 21900 | 0.002037 | 9.303e-08 | 0.002037 | 9.303e-08 | ~:0(<method 'tell' of '_io.BytesIO' objects>) |
| 494 | 0.002028 | 4.106e-06 | 0.005661 | 1.146e-05 | core.py:1571(__init__) |
| 1110 | 0.002025 | 1.825e-06 | 0.03809 | 3.432e-05 | mlir.py:442(__init__) |
| 2385 | 0.00202 | 8.47e-07 | 0.003439 | 1.442e-06 | source_info_util.py:93(__add__) |
| 746 | 0.002011 | 2.696e-06 | 0.005008 | 6.713e-06 | core.py:1890(_dim_handler_and_canonical) |
| 2012 | 0.00198 | 9.841e-07 | 0.01264 | 6.28e-06 | api_util.py:622(<genexpr>) |
| 1200 | 0.001975 | 1.645e-06 | 0.008841 | 7.368e-06 | pxla.py:1695(__call__) |
| 545/46 | 0.001973 | 4.289e-05 | 0.7005 | 0.01523 | pjit.py:916(_pjit_jaxpr) |
| 6749 | 0.001945 | 2.882e-07 | 0.001945 | 2.882e-07 | partial_eval.py:1528(_assert_live) |
| 9469 | 0.00193 | 2.038e-07 | 0.002555 | 2.699e-07 | pxla.py:520(_is_unspecified) |
| 64 | 0.001911 | 2.986e-05 | 0.00868 | 0.0001356 | pxla.py:2834(_get_input_indices) |
| 21300 | 0.001907 | 8.953e-08 | 0.001907 | 8.953e-08 | pickle.py:605(persistent_id) |
| 1409 | 0.00188 | 1.334e-06 | 0.003149 | 2.235e-06 | inspect.py:735(unwrap) |
| 2699 | 0.001866 | 6.914e-07 | 0.01884 | 6.981e-06 | api_util.py:557(shaped_abstractify) |
| 1090 | 0.001855 | 1.702e-06 | 0.003757 | 3.446e-06 | pjit.py:738(hashable_pytree) |
| 2131 | 0.001835 | 8.612e-07 | 0.002523 | 1.184e-06 | core.py:2506(<listcomp>) |
| 300 | 0.001816 | 6.055e-06 | 0.004579 | 1.526e-05 | numpy_pickle.py:286(__init__) |
| 3004 | 0.00181 | 6.024e-07 | 0.00181 | 6.024e-07 | <string>:2(__init__) |
| 12151 | 0.001809 | 1.489e-07 | 0.001809 | 1.489e-07 | ~:0(<method 'startswith' of 'str' objects>) |
| 10800 | 0.00178 | 1.648e-07 | 0.00178 | 1.648e-07 | pickle.py:514(put) |
| 1200 | 0.001774 | 1.478e-06 | 0.001774 | 1.478e-06 | ~:0(<built-in method builtins.__import__>) |
| 300 | 0.001773 | 5.909e-06 | 0.1754 | 0.0005846 | pickle.py:476(dump) |
| 250 | 0.001764 | 7.055e-06 | 0.002618 | 1.047e-05 | partial_eval.py:1696(<setcomp>) |
| 300 | 0.001753 | 5.842e-06 | 0.001956 | 6.519e-06 | pickle.py:409(__init__) |
| 328 | 0.001739 | 5.303e-06 | 0.016 | 4.877e-05 | mlir.py:1291(broadcast_in_dim) |
| 300 | 0.001734 | 5.78e-06 | 0.002626 | 8.754e-06 | <frozen posixpath>:71(join) |
| 102 | 0.001721 | 1.688e-05 | 0.003919 | 3.842e-05 | _stablehlo_ops_gen.py:1660(__init__) |
| 2739 | 0.001719 | 6.275e-07 | 0.006239 | 2.278e-06 | partial_eval.py:1687(<lambda>) |
| 704 | 0.001714 | 2.434e-06 | 0.003585 | 5.092e-06 | api_util.py:600(fun_sourceinfo) |
| 4728 | 0.00171 | 3.617e-07 | 0.00171 | 3.617e-07 | core.py:903(__eq__) |
| 10399 | 0.001707 | 1.642e-07 | 0.001707 | 1.642e-07 | ~:0(<method 'endswith' of 'str' objects>) |
| 300 | 0.001704 | 5.681e-06 | 0.2207 | 0.0007357 | logx.py:180(save_state) |
| 545 | 0.001692 | 3.105e-06 | 0.04972 | 9.123e-05 | api_util.py:590(debug_info) |
| 10274 | 0.001691 | 1.646e-07 | 0.001691 | 1.646e-07 | fromnumeric.py:2328(_any_dispatcher) |
| 2676 | 0.001679 | 6.275e-07 | 0.002243 | 8.381e-07 | enum.py:688(__call__) |
| 1595 | 0.001669 | 1.046e-06 | 0.00958 | 6.006e-06 | tree_util.py:698(generate_key_paths) |
| 1793 | 0.001641 | 9.15e-07 | 0.001976 | 1.102e-06 | core.py:2138(join_named_shapes) |
| 503 | 0.001607 | 3.194e-06 | 0.123 | 0.0002446 | pjit.py:1383(pjit_staging_rule) |
| 1881 | 0.001604 | 8.527e-07 | 0.003662 | 1.947e-06 | partial_eval.py:1798(<listcomp>) |
| 1686 | 0.001601 | 9.499e-07 | 0.004384 | 2.6e-06 | core.py:3018(<genexpr>) |
| 956 | 0.001557 | 1.629e-06 | 0.004284 | 4.481e-06 | lax.py:1563(_naryop_weak_type_rule) |
| 159/30 | 0.001551 | 5.168e-05 | 0.6968 | 0.02323 | pjit.py:864(_create_pjit_jaxpr) |
| 168 | 0.001548 | 9.214e-06 | 0.01823 | 0.0001085 | lax.py:2370(_dot_general_shape_rule) |
| 1289 | 0.001494 | 1.159e-06 | 0.001587 | 1.231e-06 | core.py:94(__init__) |
| 551/541 | 0.001488 | 2.75e-06 | 3.898 | 0.007205 | core.py:2586(bind) |
| 2186 | 0.001462 | 6.69e-07 | 0.02742 | 1.254e-05 | mlir.py:1112(read) |
| 64 | 0.001443 | 2.254e-05 | 0.01708 | 0.0002669 | partial_eval.py:1347(dce_jaxpr_consts) |
| 4765 | 0.001432 | 3.004e-07 | 0.001777 | 3.728e-07 | utils.py:58(<genexpr>) |
| 488 | 0.001427 | 2.925e-06 | 0.003649 | 7.477e-06 | dispatch.py:263(log_elapsed_time) |
| 971 | 0.001426 | 1.468e-06 | 0.006978 | 7.186e-06 | core.py:3016(str_eqn_compact) |
| 1232 | 0.001423 | 1.155e-06 | 0.001423 | 1.155e-06 | pxla.py:430(<listcomp>) |
| 801 | 0.001409 | 1.759e-06 | 0.00193 | 2.409e-06 | inspect.py:167(get_annotations) |
| 1200 | 0.001385 | 1.154e-06 | 0.001986 | 1.655e-06 | pickle.py:322(_getattribute) |
| 300 | 0.001384 | 4.613e-06 | 0.001962 | 6.541e-06 | logx.py:303(<listcomp>) |
| 1787 | 0.001383 | 7.738e-07 | 0.002677 | 1.498e-06 | core.py:1944(symbolic_equal_shape) |
| 77 | 0.001382 | 1.795e-05 | 0.001382 | 1.795e-05 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.get>) |
| 2746 | 0.001366 | 4.973e-07 | 0.001366 | 4.973e-07 | lax.py:1564(<genexpr>) |
| 47 | 0.001365 | 2.904e-05 | 0.01372 | 0.0002919 | lax.py:3684(_unary_reduce_lower) |
| 328 | 0.001342 | 4.092e-06 | 0.006921 | 2.11e-05 | _stablehlo_ops_gen.py:699(__init__) |
| 62 | 0.001329 | 2.144e-05 | 4.131 | 0.06663 | dispatch.py:485(compile_or_get_cached) |
| 44 | 0.001325 | 3.012e-05 | 0.005321 | 0.0001209 | pjit.py:1130(_resolve_in_shardings) |
| 1857 | 0.001315 | 7.079e-07 | 0.001833 | 9.873e-07 | core.py:150(jaxprs_in_params) |
| 954 | 0.001308 | 1.371e-06 | 0.001644 | 1.723e-06 | util.py:80(unzip2) |
| 3300 | 0.001306 | 3.958e-07 | 0.001902 | 5.763e-07 | pickle.py:526(get) |
| 2393 | 0.001305 | 5.455e-07 | 0.003338 | 1.395e-06 | core.py:414(write) |
| 558 | 0.001303 | 2.335e-06 | 0.001303 | 2.335e-06 | mlir.py:557(sharded_aval) |
| 1048 | 0.00127 | 1.212e-06 | 0.01342 | 1.281e-05 | mlir.py:494(replace) |
| 6662 | 0.001267 | 1.902e-07 | 0.001267 | 1.902e-07 | ~:0(<method 'encode' of 'str' objects>) |
| 1721 | 0.001263 | 7.341e-07 | 0.002899 | 1.685e-06 | tree_util.py:547(keystr) |
| 227/90 | 0.001256 | 1.395e-05 | 0.7132 | 0.007924 | partial_eval.py:2041(trace_to_jaxpr_dynamic) |
| 3600 | 0.001254 | 3.484e-07 | 0.002497 | 6.936e-07 | pickle.py:740(save_none) |
| 1051 | 0.001251 | 1.19e-06 | 0.005532 | 5.264e-06 | lax.py:4547(<setcomp>) |
| 1216 | 0.001238 | 1.018e-06 | 0.003833 | 3.152e-06 | xla.py:175(_canonicalize_ndarray_dtype) |
| 470 | 0.001236 | 2.63e-06 | 0.007969 | 1.695e-05 | core.py:1082(new_main) |
| 410 | 0.001231 | 3.003e-06 | 0.001235 | 3.013e-06 | ~:0(<method 'format' of 'str' objects>) |
| 1078 | 0.001209 | 1.121e-06 | 0.001383 | 1.283e-06 | mlir.py:809(__init__) |
| 3477 | 0.001205 | 3.464e-07 | 0.001525 | 4.387e-07 | inspect.py:3018(<genexpr>) |
| 2062 | 0.001194 | 5.789e-07 | 0.002518 | 1.221e-06 | pjit.py:82(_is_unspecified_or_auto) |
| 96 | 0.001184 | 1.234e-05 | 0.005395 | 5.62e-05 | inspect.py:1336(getfullargspec) |
| 2339 | 0.001175 | 5.026e-07 | 0.002928 | 1.252e-06 | util.py:423(__hash__) |
| 545 | 0.001173 | 2.152e-06 | 0.001394 | 2.558e-06 | sharding.py:79(<setcomp>) |
| 5036 | 0.001166 | 2.314e-07 | 0.002092 | 4.155e-07 | core.py:1932(<genexpr>) |
| 3354 | 0.001158 | 3.454e-07 | 0.001684 | 5.02e-07 | lax.py:4539(<genexpr>) |
| 1099 | 0.001157 | 1.053e-06 | 0.002634 | 2.397e-06 | linear_util.py:221(transformation_with_aux) |
| 398 | 0.001152 | 2.895e-06 | 0.008577 | 2.155e-05 | lax.py:4559(_check_shapelike) |
| 1655 | 0.001148 | 6.935e-07 | 0.002266 | 1.369e-06 | partial_eval.py:1706(<listcomp>) |
| 840 | 0.001142 | 1.359e-06 | 0.001142 | 1.359e-06 | sharding_impls.py:644(_to_xla_op_sharding) |
| 1200 | 0.00114 | 9.503e-07 | 0.03655 | 3.046e-05 | pickle.py:1122(save_type) |
| 1294 | 0.001122 | 8.67e-07 | 0.001466 | 1.133e-06 | linear_util.py:141(wrap) |
| 494 | 0.001119 | 2.264e-06 | 0.008577 | 1.736e-05 | abstract_arrays.py:52(canonical_concrete_aval) |
| 1736 | 0.00111 | 6.395e-07 | 0.00111 | 6.395e-07 | ~:0(<slot wrapper '__init__' of 'object' objects>) |
| 971 | 0.001103 | 1.136e-06 | 0.002356 | 2.426e-06 | mlir.py:830(subset) |
| 165 | 0.001099 | 6.664e-06 | 0.03134 | 0.0001899 | core.py:2544(do_subst_axis_names_jaxpr) |
| 77 | 0.001068 | 1.387e-05 | 0.00248 | 3.22e-05 | _func_ops_ext.py:93(add_entry_block) |
| 2263 | 0.001066 | 4.713e-07 | 0.00424 | 1.874e-06 | core.py:1930(is_constant_shape) |
| 160 | 0.001065 | 6.656e-06 | 0.009239 | 5.774e-05 | lax.py:2639(_broadcast_in_dim_shape_rule) |
| 7244 | 0.001049 | 1.448e-07 | 0.001049 | 1.448e-07 | reductions.py:188(force) |
| 500 | 0.001041 | 2.082e-06 | 0.006661 | 1.332e-05 | core.py:933(_update_thread_local_jit_state) |
| 4568 | 0.001029 | 2.253e-07 | 0.001413 | 3.094e-07 | pxla.py:511(is_auto) |
| 2943 | 0.001019 | 3.462e-07 | 0.003359 | 1.141e-06 | mlir.py:1119(aval) |
| 1200 | 0.001017 | 8.471e-07 | 0.001017 | 8.471e-07 | pxla.py:3259(<listcomp>) |
| 1232 | 0.001016 | 8.244e-07 | 0.001352 | 1.097e-06 | pxla.py:451(<listcomp>) |
| 1496 | 0.001012 | 6.764e-07 | 0.002402 | 1.606e-06 | core.py:1883(_get_special_dim_handler) |
| 2945 | 0.001001 | 3.4e-07 | 0.001296 | 4.401e-07 | core.py:1509(<lambda>) |
| 844 | 0.0009988 | 1.183e-06 | 0.0009988 | 1.183e-06 | ~:0(<built-in method jaxlib.xla_extension.from_proto>) |
| 2400 | 0.0009848 | 4.103e-07 | 0.0009848 | 4.103e-07 | pxla.py:3256(<genexpr>) |
| 2599 | 0.0009384 | 3.61e-07 | 0.001372 | 5.277e-07 | core.py:411(read) |
| 250 | 0.0009361 | 3.744e-06 | 0.01384 | 5.536e-05 | partial_eval.py:2310(_input_type_to_tracers) |
| 831 | 0.0009311 | 1.12e-06 | 0.0009311 | 1.12e-06 | partial_eval.py:1538(make_jaxpr_effects) |
| 386 | 0.000923 | 2.391e-06 | 0.001584 | 4.104e-06 | linear_util.py:146(populate_stores) |
| 6625 | 0.0009201 | 1.389e-07 | 0.0009201 | 1.389e-07 | ~:0(<method 'add' of 'set' objects>) |
| 96/16 | 0.0009188 | 5.742e-05 | 0.08514 | 0.005321 | module.py:399(wrapped) |
| 108 | 0.0009161 | 8.482e-06 | 0.003483 | 3.225e-05 | pxla.py:2334(_get_and_check_device_assignment) |
| 4/2 | 0.0009121 | 0.0004561 | 0.03021 | 0.0151 | ad.py:181(backward_pass) |
| 4200 | 0.0009051 | 2.155e-07 | 0.0009051 | 2.155e-07 | core.py:117(reset) |
| 551 | 0.0008917 | 1.618e-06 | 0.04724 | 8.574e-05 | core.py:2490(used_axis_names) |
| 321 | 0.0008794 | 2.74e-06 | 0.009339 | 2.909e-05 | partial_eval.py:1730(new_const) |
| 6000 | 0.0008781 | 1.463e-07 | 0.0008781 | 1.463e-07 | ~:0(<method 'discard' of 'set' objects>) |
| 672 | 0.0008697 | 1.294e-06 | 0.00553 | 8.23e-06 | core.py:1948(greater_equal_dim) |
| 600 | 0.0008648 | 1.441e-06 | 0.1237 | 0.0002062 | lax.py:4254(_array_copy) |
| 3014 | 0.0008598 | 2.853e-07 | 0.0008598 | 2.853e-07 | source_info_util.py:57(wrap) |
| 1033 | 0.0008555 | 8.282e-07 | 0.001107 | 1.071e-06 | effects.py:59(filter_in) |
| 423 | 0.0008514 | 2.013e-06 | 0.00156 | 3.688e-06 | pxla.py:3369(is_op_sharding_replicated) |
| 305 | 0.0008501 | 2.787e-06 | 0.001519 | 4.982e-06 | source_info_util.py:133(_raw_frame_to_frame) |
| 3249 | 0.0008427 | 2.594e-07 | 0.001073 | 3.301e-07 | core.py:1919(is_constant_dim) |
| 1881 | 0.0008393 | 4.462e-07 | 0.001081 | 5.748e-07 | partial_eval.py:1636(<listcomp>) |
| 1881 | 0.0008209 | 4.364e-07 | 0.0008209 | 4.364e-07 | partial_eval.py:1794(<listcomp>) |
| 6000 | 0.0008172 | 1.362e-07 | 0.0008172 | 1.362e-07 | ~:0(<built-in method _thread.get_ident>) |
| 1515 | 0.0008167 | 5.391e-07 | 0.004823 | 3.183e-06 | dispatch.py:399(check_arg) |
| 500 | 0.0008145 | 1.629e-06 | 0.003533 | 7.067e-06 | source_info_util.py:255(reset_name_stack) |
| 1580 | 0.0008123 | 5.141e-07 | 0.001068 | 6.763e-07 | core.py:2535(<listcomp>) |
| 112 | 0.0008097 | 7.23e-06 | 0.01399 | 0.0001249 | partial_eval.py:215(default_process_primitive) |
| 1397 | 0.0008097 | 5.796e-07 | 0.01397 | 1e-05 | mlir.py:159(aval_to_ir_types) |
| 599 | 0.0008031 | 1.341e-06 | 0.002655 | 4.433e-06 | dtypes.py:157(_scalar_type_to_dtype) |
| 165 | 0.000803 | 4.867e-06 | 0.03221 | 0.0001952 | core.py:2559(used_axis_names_jaxpr) |
| 44 | 0.0007988 | 1.816e-05 | 0.3509 | 0.007974 | pjit.py:1307(_pjit_lower) |
| 190 | 0.0007934 | 4.176e-06 | 0.002662 | 1.401e-05 | partial_eval.py:1741(_new_const) |
| 4765 | 0.0007931 | 1.664e-07 | 0.0007931 | 1.664e-07 | utils.py:115(<genexpr>) |
| 5079 | 0.0007837 | 1.543e-07 | 0.0007837 | 1.543e-07 | dtypes.py:637(<genexpr>) |
| 4144 | 0.0007833 | 1.89e-07 | 0.0007833 | 1.89e-07 | ~:0(<built-in method builtins.iter>) |
| 1500 | 0.0007817 | 5.212e-07 | 0.001316 | 8.776e-07 | pickle.py:874(save_tuple) |
| 1721 | 0.0007796 | 4.53e-07 | 0.001411 | 8.201e-07 | tree_util.py:548(<listcomp>) |
| 1090 | 0.0007786 | 7.143e-07 | 0.003546 | 3.253e-06 | pjit.py:752(_create_sharding_for_array) |
| 277 | 0.0007717 | 2.786e-06 | 0.006333 | 2.286e-05 | dtypes.py:649(result_type) |
| 67 | 0.0007695 | 1.148e-05 | 4.148 | 0.06191 | pxla.py:2808(compile) |
| 299 | 0.0007635 | 2.554e-06 | 0.03204 | 0.0001072 | mlir.py:202(ir_constants) |
| 2036 | 0.0007618 | 3.741e-07 | 0.00129 | 6.337e-07 | partial_eval.py:1366(write) |
| 4040 | 0.000759 | 1.879e-07 | 0.000759 | 1.879e-07 | core.py:87(<lambda>) |
| 77/64 | 0.0007587 | 1.186e-05 | 0.01201 | 0.0001877 | partial_eval.py:1339(dce_jaxpr) |
| 494 | 0.0007555 | 1.529e-06 | 0.0007555 | 1.529e-06 | sharding_impls.py:647(get_replicated) |
| 494 | 0.0007487 | 1.516e-06 | 0.01294 | 2.619e-05 | core.py:1344(concrete_aval) |
| 1500 | 0.0007479 | 4.986e-07 | 0.001267 | 8.444e-07 | pickle.py:744(save_bool) |
| 971 | 0.0007433 | 7.655e-07 | 0.002133 | 2.197e-06 | xla.py:63(get_canonical_source_file) |
| 1642 | 0.0007345 | 4.473e-07 | 0.0007345 | 4.473e-07 | mlir.py:1125(write) |
| 1881 | 0.0007292 | 3.877e-07 | 0.0009572 | 5.089e-07 | partial_eval.py:1576(add_eqn) |
| 3931 | 0.0007246 | 1.843e-07 | 0.0007246 | 1.843e-07 | core.py:89(<lambda>) |
| 545 | 0.0007201 | 1.321e-06 | 0.001189 | 2.182e-06 | api_util.py:288(argnums_partial_except) |
| 32 | 0.0007193 | 2.248e-05 | 0.05222 | 0.001632 | basic.py:158(__call__) |
| 1061 | 0.0007191 | 6.778e-07 | 0.002086 | 1.966e-06 | utils.py:38(<lambda>) |
| 1580 | 0.0007115 | 4.503e-07 | 0.01229 | 7.779e-06 | core.py:2537(<listcomp>) |
| 347 | 0.0007098 | 2.046e-06 | 0.009128 | 2.63e-05 | abstract_arrays.py:74(_make_concrete_python_scalar) |
| 705 | 0.0007064 | 1.002e-06 | 0.0267 | 3.787e-05 | inspect.py:3023(from_callable) |
| 2840 | 0.0007045 | 2.48e-07 | 0.0007163 | 2.522e-07 | ~:0(<built-in method from_iterable>) |
| 547/44 | 0.0007024 | 1.596e-05 | 4.273 | 0.09712 | traceback_util.py:162(reraise_with_filtered_traceback) |
| 1655 | 0.0007003 | 4.232e-07 | 0.005855 | 3.537e-06 | partial_eval.py:1707(<listcomp>) |
| 2715 | 0.0006938 | 2.555e-07 | 0.0006938 | 2.555e-07 | dtypes.py:593(check_valid_dtype) |
| 3386 | 0.0006923 | 2.045e-07 | 0.000965 | 2.85e-07 | partial_eval.py:1363(read) |
| 26 | 0.0006898 | 2.653e-05 | 0.001965 | 7.558e-05 | _stablehlo_ops_gen.py:1156(__init__) |
| 3600 | 0.000688 | 1.911e-07 | 0.000688 | 1.911e-07 | ~:0(<method 'setdefault' of 'dict' objects>) |
| 62 | 0.0006878 | 1.109e-05 | 0.0006878 | 1.109e-05 | pxla.py:2874(<listcomp>) |
| 289/274 | 0.0006827 | 2.492e-06 | 0.003219 | 1.175e-05 | dispatch.py:325(jaxpr_shardings) |
| 3766 | 0.0006818 | 1.81e-07 | 0.0006818 | 1.81e-07 | core.py:88(<lambda>) |
| 1916 | 0.0006762 | 3.529e-07 | 0.000821 | 4.285e-07 | mlir.py:544(_unwrap_singleton_ir_values) |
| 40 | 0.0006743 | 1.686e-05 | 0.001545 | 3.863e-05 | _stablehlo_ops_gen.py:4890(__init__) |
| 305 | 0.0006735 | 2.208e-06 | 0.001345 | 4.408e-06 | source_info_util.py:147(user_frames) |
| 500 | 0.0006733 | 1.347e-06 | 0.00164 | 3.28e-06 | source_info_util.py:245(set_name_stack) |
| 3405 | 0.0006683 | 1.963e-07 | 0.0009963 | 2.926e-07 | partial_eval.py:1691(vars_in_shape) |
| 705 | 0.0006657 | 9.442e-07 | 0.02736 | 3.882e-05 | inspect.py:3277(signature) |
| 1934 | 0.0006631 | 3.429e-07 | 0.0006631 | 3.429e-07 | linear_util.py:129(__init__) |
| 2012 | 0.0006607 | 3.284e-07 | 0.009808 | 4.875e-06 | pjit.py:475(<genexpr>) |
| 62 | 0.0006589 | 1.063e-05 | 0.00168 | 2.71e-05 | pxla.py:2862(get_gspmd_shardings_from_executable) |
| 225 | 0.000658 | 2.925e-06 | 0.001238 | 5.502e-06 | core.py:137(replace) |
| 159 | 0.0006556 | 4.123e-06 | 0.0009484 | 5.965e-06 | api_util.py:116(flattened_fun_in_tree) |
| 62 | 0.0006542 | 1.055e-05 | 0.03529 | 0.0005692 | mlir.py:340(make_ir_context) |
| 2748 | 0.0006527 | 2.375e-07 | 0.0009383 | 3.415e-07 | partial_eval.py:2330(_substitute_vars_in_type) |
| 640 | 0.0006508 | 1.017e-06 | 0.001032 | 1.612e-06 | linear_util.py:234(wrap_init) |
| 50 | 0.0006478 | 1.296e-05 | 0.0006478 | 1.296e-05 | ~:0(<method 'argsort' of 'numpy.ndarray' objects>) |
| 865 | 0.0006467 | 7.476e-07 | 0.004981 | 5.758e-06 | lax.py:4581(<genexpr>) |
| 165 | 0.0006413 | 3.886e-06 | 0.002491 | 1.51e-05 | core.py:160(subjaxprs) |
| 336 | 0.00064 | 1.905e-06 | 0.000841 | 2.503e-06 | lax.py:2441(tuple_delete) |
| 305 | 0.0006329 | 2.075e-06 | 0.0006329 | 2.075e-06 | ~:0(<built-in method jaxlib.xla_extension.raw_frames>) |
| 3206 | 0.0006254 | 1.951e-07 | 0.0006254 | 1.951e-07 | ~:0(<method 'pop' of 'list' objects>) |
| 600 | 0.0006177 | 1.03e-06 | 0.001028 | 1.713e-06 | pickle.py:751(save_long) |
| 32 | 0.0006166 | 1.927e-05 | 0.006745 | 0.0002108 | util.py:400(_broadcast_to) |
| 1485 | 0.0006138 | 4.133e-07 | 0.0009181 | 6.183e-07 | linear_util.py:101(val) |
| 547 | 0.0006065 | 1.109e-06 | 0.0008078 | 1.477e-06 | weakref.py:507(setdefault) |
| 730 | 0.0006022 | 8.249e-07 | 0.001358 | 1.86e-06 | core.py:3008(_compact_eqn_should_include) |
| 924/567 | 0.0006009 | 1.06e-06 | 0.001287 | 2.27e-06 | core.py:2245(primitive_uses_outfeed) |
| 5922 | 0.0005969 | 1.008e-07 | 0.0005969 | 1.008e-07 | inspect.py:2736(name) |
| 55 | 0.0005953 | 1.082e-05 | 0.0008793 | 1.599e-05 | function_base.py:5054(delete) |
| 1223 | 0.0005899 | 4.824e-07 | 0.0009024 | 7.378e-07 | mlir.py:545(wrap_singleton_ir_values) |
| 290 | 0.0005871 | 2.025e-06 | 0.01904 | 6.566e-05 | source_info_util.py:160(user_frame) |
| 504 | 0.000586 | 1.163e-06 | 0.000586 | 1.163e-06 | core.py:287(gensym) |
| 2048 | 0.0005847 | 2.855e-07 | 0.0005847 | 2.855e-07 | ~:0(<built-in method builtins.repr>) |
| 1998 | 0.0005842 | 2.924e-07 | 0.0006639 | 3.323e-07 | mlir.py:1195(<genexpr>) |
| 165 | 0.0005833 | 3.535e-06 | 0.02322 | 0.0001408 | core.py:2552(<listcomp>) |
| 1982 | 0.0005813 | 2.933e-07 | 0.002627 | 1.326e-06 | lax.py:1547(<genexpr>) |
| 189 | 0.0005798 | 3.067e-06 | 0.0009007 | 4.765e-06 | source_info_util.py:73(extend) |
| 1051/1027 | 0.000577 | 5.618e-07 | 0.01427 | 1.39e-05 | core.py:1354(get_aval) |
| 62 | 0.000577 | 9.306e-06 | 0.001507 | 2.431e-05 | function_base.py:2400(_vectorize_call) |
| 728 | 0.0005761 | 7.913e-07 | 0.002042 | 2.805e-06 | core.py:1236(same_referent) |
| 84 | 0.0005759 | 6.856e-06 | 0.0008647 | 1.029e-05 | _dtype.py:330(_name_includes_bit_suffix) |
| 77 | 0.0005724 | 7.434e-06 | 0.0005724 | 7.434e-06 | _func_ops_ext.py:75(type) |
| 1998 | 0.0005714 | 2.86e-07 | 0.0007008 | 3.508e-07 | mlir.py:1196(<genexpr>) |
| 895 | 0.0005711 | 6.381e-07 | 0.002335 | 2.609e-06 | tree_util.py:64(tree_unflatten) |
| 44 | 0.0005686 | 1.292e-05 | 3.545 | 0.08058 | pjit.py:1219(_pjit_call_impl) |
| 4 | 0.0005657 | 0.0001414 | 0.004531 | 0.001133 | partial_eval.py:876(tracers_to_jaxpr) |
| 2676 | 0.0005637 | 2.106e-07 | 0.0005637 | 2.106e-07 | enum.py:1093(__new__) |
| 4 | 0.0005563 | 0.0001391 | 0.00198 | 0.0004951 | util.py:176(toposort) |
| 214 | 0.0005494 | 2.567e-06 | 0.001395 | 6.52e-06 | lax.py:2226(_convert_elt_type_folding_rule) |
| 417 | 0.0005406 | 1.296e-06 | 0.002959 | 7.097e-06 | mlir.py:67(dense_int_elements) |
| 477 | 0.0005384 | 1.129e-06 | 0.0006306 | 1.322e-06 | util.py:142(concatenate) |
| 474 | 0.0005369 | 1.133e-06 | 0.0008888 | 1.875e-06 | utils.py:117(_standard_weak_type_rule) |
| 48 | 0.0005329 | 1.11e-05 | 0.00791 | 0.0001648 | utils.py:41(auto_repr) |
| 62 | 0.0005312 | 8.567e-06 | 0.0005312 | 8.567e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.create>) |
| 300 | 0.0005295 | 1.765e-06 | 0.002265 | 7.551e-06 | pickle.py:212(end_framing) |
| 1530 | 0.0005266 | 3.442e-07 | 0.0005266 | 3.442e-07 | source_info_util.py:231(current_name_stack) |
| 32 | 0.0005224 | 1.632e-05 | 0.001182 | 3.694e-05 | _stablehlo_ops_gen.py:4858(__init__) |
| 192 | 0.0005208 | 2.713e-06 | 0.001864 | 9.707e-06 | base.py:146(module) |
| 268 | 0.0005201 | 1.941e-06 | 0.0005201 | 1.941e-06 | getlimits.py:668(__init__) |
| 15 | 0.0005174 | 3.449e-05 | 0.1197 | 0.007981 | pjit.py:1437(_pjit_lowering) |
| 764 | 0.0005164 | 6.759e-07 | 0.0006396 | 8.372e-07 | lax.py:1541(<setcomp>) |
| 33 | 0.0005162 | 1.564e-05 | 0.001955 | 5.925e-05 | lax.py:1020(_get_monoid_reducer) |
| 971 | 0.000516 | 5.314e-07 | 0.0007855 | 8.09e-07 | <string>:2(__hash__) |
| 168 | 0.0005158 | 3.07e-06 | 0.001475 | 8.781e-06 | lax.py:2432(_dot_general_shape_computation) |
| 567 | 0.0005048 | 8.904e-07 | 0.007197 | 1.269e-05 | mlir.py:174(aval_to_ir_type) |
| 77 | 0.0005043 | 6.549e-06 | 0.0005815 | 7.552e-06 | _func_ops_ext.py:67(is_external) |
| 1417 | 0.0005009 | 3.535e-07 | 0.0008348 | 5.891e-07 | pjit.py:1385(<genexpr>) |
| 66 | 0.0004995 | 7.568e-06 | 0.0009705 | 1.471e-05 | partial_eval.py:959(convert_constvars_jaxpr) |
| 1050 | 0.0004992 | 4.754e-07 | 0.0007097 | 6.759e-07 | dataclasses.py:1244(_is_dataclass_instance) |
| 545 | 0.0004976 | 9.13e-07 | 0.003094 | 5.676e-06 | sharding.py:76(addressable_devices) |
| 57 | 0.0004953 | 8.689e-06 | 0.005807 | 0.0001019 | lax.py:2605(_dot_general_lower) |
| 419 | 0.0004951 | 1.182e-06 | 0.001732 | 4.133e-06 | mlir.py:1515(get_sharding_attr) |
| 840 | 0.0004899 | 5.833e-07 | 0.0004899 | 5.833e-07 | lax.py:1538(<listcomp>) |
| 788 | 0.0004888 | 6.203e-07 | 0.0004888 | 6.203e-07 | ~:0(<built-in method time.time>) |
| 250 | 0.000487 | 1.948e-06 | 0.0008007 | 3.203e-06 | partial_eval.py:1565(__init__) |
| 33 | 0.0004815 | 1.459e-05 | 0.01628 | 0.0004933 | reductions.py:71(_reduction) |
| 411 | 0.000478 | 1.163e-06 | 0.0005956 | 1.449e-06 | core.py:178(__init__) |
| 250 | 0.0004759 | 1.903e-06 | 0.001614 | 6.455e-06 | partial_eval.py:1681(<dictcomp>) |
| 62 | 0.0004728 | 7.627e-06 | 0.04226 | 0.0006817 | mlir.py:1212(f_lowered) |
| 626 | 0.000467 | 7.46e-07 | 0.0004907 | 7.839e-07 | api_util.py:273(_ensure_inbounds) |
| 1596 | 0.0004657 | 2.918e-07 | 0.001601 | 1.003e-06 | core.py:1233(get_referent) |
| 62 | 0.0004599 | 7.417e-06 | 0.01158 | 0.0001868 | pxla.py:2928(build_unsafe_call) |
| 27 | 0.0004586 | 1.698e-05 | 0.001045 | 3.87e-05 | _stablehlo_ops_gen.py:4927(__init__) |
| 93 | 0.0004553 | 4.895e-06 | 0.002483 | 2.67e-05 | pjit.py:842(_process_in_axis_resources) |
| 141 | 0.0004547 | 3.225e-06 | 0.001186 | 8.413e-06 | array.py:720(_array_global_result_handler) |
| 681 | 0.0004541 | 6.668e-07 | 0.0004541 | 6.668e-07 | <string>:2(__eq__) |
| 64 | 0.0004487 | 7.01e-06 | 0.001301 | 2.033e-05 | base.py:464(get_parameter) |
| 363 | 0.0004463 | 1.229e-06 | 0.002655 | 7.314e-06 | tree_util.py:580(flatten_with_keys) |
| 1200 | 0.000445 | 3.708e-07 | 0.000722 | 6.017e-07 | pickle.py:335(whichmodule) |
| 1604 | 0.0004443 | 2.77e-07 | 0.0005704 | 3.556e-07 | inspect.py:378(isfunction) |
| 160 | 0.000442 | 2.763e-06 | 0.01075 | 6.719e-05 | lax.py:2804(_broadcast_in_dim_abstract_eval) |
| 154 | 0.0004394 | 2.853e-06 | 0.0009825 | 6.38e-06 | util.py:155(<listcomp>) |
| 108 | 0.0004377 | 4.053e-06 | 0.0007541 | 6.982e-06 | api_util.py:326(argnames_partial_except) |
| 235 | 0.0004263 | 1.814e-06 | 0.00796 | 3.387e-05 | api_util.py:538(_shaped_abstractify_slow) |
| 62 | 0.0004231 | 6.824e-06 | 0.00479 | 7.725e-05 | mlir.py:770(_set_up_aliases) |
| 2180 | 0.0004218 | 1.935e-07 | 0.0004218 | 1.935e-07 | util.py:414(__init__) |
| 335 | 0.0004211 | 1.257e-06 | 0.0007794 | 2.327e-06 | core.py:184(in_avals) |
| 305 | 0.000421 | 1.38e-06 | 0.000421 | 1.38e-06 | ~:0(<built-in method jaxlib.xla_extension.code_addr2location>) |
| 159 | 0.0004207 | 2.646e-06 | 0.001487 | 9.352e-06 | api_util.py:631(jaxpr_debug_info) |
| 1296 | 0.0004197 | 3.239e-07 | 0.0004197 | 3.239e-07 | ~:0(<method 'split' of 'str' objects>) |
| 173 | 0.0004177 | 2.414e-06 | 0.0009601 | 5.55e-06 | util.py:299(__hash__) |
| 77 | 0.0004125 | 5.356e-06 | 0.002268 | 2.945e-05 | _func_ops_gen.py:116(__init__) |
| 545 | 0.0004123 | 7.565e-07 | 0.0004123 | 7.565e-07 | api_util.py:616(<listcomp>) |
| 2676 | 0.0004108 | 1.535e-07 | 0.0004108 | 1.535e-07 | ~:0(<method 'isidentifier' of 'str' objects>) |
| 399 | 0.0004102 | 1.028e-06 | 0.0008545 | 2.142e-06 | partial_eval.py:676(__init__) |
| 346 | 0.0004088 | 1.181e-06 | 0.0006669 | 1.928e-06 | core.py:188(out_avals) |
| 842 | 0.0004055 | 4.816e-07 | 0.0005666 | 6.729e-07 | pxla.py:2850(<genexpr>) |
| 300 | 0.0004052 | 1.351e-06 | 0.0004052 | 1.351e-06 | pickle.py:209(start_framing) |
| 408 | 0.0004002 | 9.809e-07 | 0.0006045 | 1.482e-06 | core.py:318(__init__) |
| 740 | 0.0003984 | 5.384e-07 | 0.0003984 | 5.384e-07 | ~:0(<built-in method jaxlib.xla_extension.process_index>) |
| 558 | 0.0003964 | 7.103e-07 | 0.003672 | 6.581e-06 | mlir.py:774(<lambda>) |
| 1251 | 0.0003961 | 3.166e-07 | 0.0009241 | 7.387e-07 | partial_eval.py:693(aval) |
| 1743 | 0.0003935 | 2.257e-07 | 0.001041 | 5.975e-07 | core.py:1934(symbolic_equal_dim) |
| 1299 | 0.0003932 | 3.027e-07 | 0.000546 | 4.203e-07 | partial_eval.py:131(get_aval) |
| 236 | 0.0003931 | 1.666e-06 | 0.001787 | 7.57e-06 | xla.py:205(_make_abstract_python_scalar) |
| 421 | 0.000393 | 9.334e-07 | 0.000393 | 9.334e-07 | sharding_impls.py:630(device_set) |
| 1008 | 0.0003887 | 3.856e-07 | 0.0003887 | 3.856e-07 | lax.py:1515(<listcomp>) |
| 900 | 0.0003878 | 4.309e-07 | 0.0003878 | 4.309e-07 | ~:0(<method 'write' of '_io.BufferedWriter' objects>) |
| 77 | 0.000384 | 4.987e-06 | 0.003375 | 4.384e-05 | _func_ops_ext.py:32(__init__) |
| 225 | 0.000384 | 1.707e-06 | 0.000504 | 2.24e-06 | api_util.py:313(_argnums_partial) |
| 498 | 0.000384 | 7.71e-07 | 0.0004548 | 9.132e-07 | core.py:855(__eq__) |
| 300 | 0.0003828 | 1.276e-06 | 0.0003828 | 1.276e-06 | ~:0(<method 'clear' of 'dict' objects>) |
| 2805 | 0.0003788 | 1.351e-07 | 0.0003788 | 1.351e-07 | ~:0(<built-in method builtins.callable>) |
| 261 | 0.0003735 | 1.431e-06 | 0.001831 | 7.016e-06 | _ods_common.py:146(get_op_results_or_values) |
| 250 | 0.0003725 | 1.49e-06 | 0.0005173 | 2.069e-06 | partial_eval.py:1698(<setcomp>) |
| 545 | 0.0003725 | 6.835e-07 | 0.006484 | 1.19e-05 | inspect.py:3206(bind) |
| 261 | 0.000371 | 1.422e-06 | 0.001331 | 5.1e-06 | _ods_common.py:161(<listcomp>) |
| 108 | 0.0003707 | 3.433e-06 | 0.001391 | 1.288e-05 | xla_bridge.py:592(local_devices) |
| 300 | 0.0003705 | 1.235e-06 | 0.0003705 | 1.235e-06 | ezpickle.py:25(__getstate__) |
| 539 | 0.0003696 | 6.856e-07 | 0.0009117 | 1.691e-06 | partial_eval.py:105(__new__) |
| 545 | 0.0003685 | 6.762e-07 | 0.0003685 | 6.762e-07 | pjit.py:219(_read_most_recent_pjit_call_executable) |
| 558 | 0.0003681 | 6.596e-07 | 0.001239 | 2.22e-06 | pxla.py:2436(<genexpr>) |
| 2676 | 0.0003662 | 1.369e-07 | 0.0003662 | 1.369e-07 | ~:0(<method '__contains__' of 'frozenset' objects>) |
| 278 | 0.0003652 | 1.314e-06 | 0.0003652 | 1.314e-06 | ~:0(<built-in method jaxlib.xla_extension.devices>) |
| 308 | 0.0003651 | 1.185e-06 | 0.0006189 | 2.009e-06 | __init__.py:1544(log) |
| 25 | 0.0003648 | 1.459e-05 | 0.009529 | 0.0003812 | lax.py:2454(_dot_general_transpose_lhs) |
| 159 | 0.0003645 | 2.292e-06 | 0.008012 | 5.039e-05 | partial_eval.py:1999(debug_info) |
| 343 | 0.0003643 | 1.062e-06 | 0.0009832 | 2.866e-06 | tree_util.py:597(<lambda>) |
| 644 | 0.0003625 | 5.628e-07 | 0.002711 | 4.21e-06 | dispatch.py:391(_valid_jaxtype) |
| 62 | 0.0003623 | 5.843e-06 | 0.0005831 | 9.404e-06 | function_base.py:2269(__init__) |
| 114 | 0.0003615 | 3.171e-06 | 0.0003615 | 3.171e-06 | lax.py:649(__str__) |
| 1857 | 0.0003611 | 1.945e-07 | 0.0003611 | 1.945e-07 | core.py:175(<lambda>) |
| 73 | 0.0003559 | 4.875e-06 | 0.002926 | 4.008e-05 | mlir.py:1462(convert_hlo) |
| 1181 | 0.000354 | 2.997e-07 | 0.0005312 | 4.498e-07 | tree_util.py:597(<genexpr>) |
| 761 | 0.0003537 | 4.647e-07 | 0.00714 | 9.382e-06 | mlir.py:898(aval_to_types) |
| 77 | 0.0003507 | 4.555e-06 | 0.0003663 | 4.757e-06 | _func_ops_ext.py:108(arg_attrs) |
| 3989 | 0.0003488 | 8.745e-08 | 0.0003488 | 8.745e-08 | inspect.py:2748(kind) |
| 224 | 0.0003454 | 1.542e-06 | 0.0004025 | 1.797e-06 | util.py:112(split_list) |
| 1200 | 0.0003444 | 2.87e-07 | 0.0003444 | 2.87e-07 | ~:0(<method 'rpartition' of 'str' objects>) |
| 93 | 0.000343 | 3.688e-06 | 0.0006267 | 6.739e-06 | lax.py:3549(_reduce_op_shape_rule) |
| 2397 | 0.0003397 | 1.417e-07 | 0.0003397 | 1.417e-07 | core.py:90(<lambda>) |
| 795 | 0.0003375 | 4.246e-07 | 0.0007249 | 9.119e-07 | partial_eval.py:697(parents) |
| 64 | 0.0003364 | 5.256e-06 | 0.001039 | 1.623e-05 | partial_eval.py:972(convert_invars_to_constvars) |
| 62 | 0.0003358 | 5.417e-06 | 0.0003358 | 5.417e-06 | ~:0(<built-in method numpy.frompyfunc>) |
| 165 | 0.0003348 | 2.029e-06 | 0.006222 | 3.771e-05 | core.py:2550(<listcomp>) |
| 48/8 | 0.0003347 | 4.184e-05 | 0.01251 | 0.001564 | module.py:110(__call__) |
| 48 | 0.000334 | 6.958e-06 | 0.002981 | 6.211e-05 | inspect.py:1522(getcallargs) |
| 41 | 0.0003298 | 8.043e-06 | 0.3483 | 0.008496 | pjit.py:1318(_pjit_lower_cached) |
| 558 | 0.0003284 | 5.885e-07 | 0.001196 | 2.143e-06 | pxla.py:2445(<genexpr>) |
| 124 | 0.0003269 | 2.636e-06 | 0.0003269 | 2.636e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.unknown>) |
| 188 | 0.0003256 | 1.732e-06 | 0.002416 | 1.285e-05 | lax.py:1493(unop_dtype_rule) |
| 62 | 0.0003223 | 5.198e-06 | 0.0003223 | 5.198e-06 | ~:0(<built-in method jaxlib.xla_extension.create>) |
| 2363 | 0.0003201 | 1.354e-07 | 0.0003201 | 1.354e-07 | core.py:2526(<genexpr>) |
| 250 | 0.0003199 | 1.28e-06 | 0.003899 | 1.56e-05 | partial_eval.py:1703(<listcomp>) |
| 77 | 0.0003198 | 4.153e-06 | 0.001878 | 2.439e-05 | _func_ops_gen.py:185(__init__) |
| 713 | 0.0003194 | 4.479e-07 | 0.0007679 | 1.077e-06 | linear_util.py:225(<lambda>) |
| 996/624 | 0.0003155 | 5.057e-07 | 0.00146 | 2.34e-06 | core.py:2233(<genexpr>) |
| 343 | 0.0003155 | 9.198e-07 | 0.001132 | 3.301e-06 | tree_util.py:240(<lambda>) |
| 1939 | 0.0003152 | 1.625e-07 | 0.0003152 | 1.625e-07 | lax.py:1574(<genexpr>) |
| 1099 | 0.0003073 | 2.796e-07 | 0.0003073 | 2.796e-07 | linear_util.py:92(store) |
| 405 | 0.0003062 | 7.561e-07 | 0.001098 | 2.71e-06 | core.py:1790(<lambda>) |
| 1032 | 0.000305 | 2.956e-07 | 0.0004603 | 4.46e-07 | partial_eval.py:2320(_substitute_tracers_in_aval) |
| 98 | 0.0003044 | 3.106e-06 | 0.0005284 | 5.392e-06 | functools.py:35(update_wrapper) |
| 1485 | 0.0003043 | 2.049e-07 | 0.0003043 | 2.049e-07 | linear_util.py:107(__nonzero__) |
| 898 | 0.0003011 | 3.353e-07 | 0.0003503 | 3.901e-07 | partial_eval.py:1404(_default_dce_rule) |
| 280 | 0.0002966 | 1.059e-06 | 0.001241 | 4.431e-06 | ad.py:186(write_cotangent) |
| 1106 | 0.0002946 | 2.663e-07 | 0.0004147 | 3.75e-07 | tree_util.py:527(__str__) |
| 84 | 0.0002924 | 3.481e-06 | 0.001273 | 1.515e-05 | _dtype.py:344(_name_get) |
| 411 | 0.0002908 | 7.076e-07 | 0.0002908 | 7.076e-07 | source_info_util.py:63(wrap) |
| 716 | 0.0002902 | 4.053e-07 | 0.0002902 | 4.053e-07 | core.py:2484(__init__) |
| 252 | 0.00029 | 1.151e-06 | 0.0009141 | 3.628e-06 | partial_eval.py:901(newvar) |
| 64 | 0.0002897 | 4.527e-06 | 0.01747 | 0.000273 | dispatch.py:369(_prune_unused_inputs) |
| 62 | 0.0002863 | 4.618e-06 | 0.00342 | 5.517e-05 | mlir.py:764(module_to_bytecode) |
| 114 | 0.0002855 | 2.505e-06 | 0.0009113 | 7.994e-06 | partial_eval.py:845(new_eqn_recipe) |
| 124 | 0.0002849 | 2.297e-06 | 0.0002849 | 2.297e-06 | ~:0(<built-in method jaxlib.xla_extension.local_devices>) |
| 704 | 0.0002835 | 4.027e-07 | 0.0004281 | 6.081e-07 | inspect.py:752(_is_wrapper) |
| 706 | 0.0002828 | 4.006e-07 | 0.0004111 | 5.823e-07 | linear_util.py:137(__name__) |
| 57 | 0.0002827 | 4.959e-06 | 0.001358 | 2.383e-05 | _stablehlo_ops_gen.py:1697(__init__) |
| 216 | 0.0002793 | 1.293e-06 | 0.0003313 | 1.534e-06 | api_util.py:347(_argnames_partial) |
| 73 | 0.0002765 | 3.788e-06 | 0.001425 | 1.952e-05 | _stablehlo_ops_gen.py:1230(__init__) |
| 500 | 0.0002745 | 5.491e-07 | 0.0002745 | 5.491e-07 | partial_eval.py:2102(extend_jaxpr_stack) |
| 70 | 0.0002734 | 3.906e-06 | 0.007525 | 0.0001075 | util.py:235(promote_shapes) |
| 62 | 0.0002712 | 4.375e-06 | 0.0004573 | 7.376e-06 | mlir.py:722(<listcomp>) |
| 2363 | 0.0002712 | 1.148e-07 | 0.0002712 | 1.148e-07 | core.py:2527(<dictcomp>) |
| 148 | 0.0002707 | 1.829e-06 | 0.0144 | 9.733e-05 | lax.py:2727(_broadcast_in_dim_staging_rule) |
| 76/2 | 0.0002688 | 0.0001344 | 0.1678 | 0.0839 | ad.py:309(process_primitive) |
| 2486 | 0.0002682 | 1.079e-07 | 0.0002682 | 1.079e-07 | partial_eval.py:1483(full_lower) |
| 274 | 0.0002661 | 9.711e-07 | 0.0003373 | 1.231e-06 | ad.py:210(read_primal) |
| 110 | 0.0002649 | 2.408e-06 | 0.001145 | 1.041e-05 | pjit.py:926(pjit_check_aval_sharding) |
| 1409 | 0.0002647 | 1.879e-07 | 0.0002647 | 1.879e-07 | ~:0(<built-in method sys.getrecursionlimit>) |
| 477 | 0.0002628 | 5.509e-07 | 0.002241 | 4.699e-06 | api_util.py:625(result_paths) |
| 545 | 0.0002627 | 4.821e-07 | 0.002325 | 4.267e-06 | pjit.py:465(<lambda>) |
| 802 | 0.0002597 | 3.238e-07 | 0.002138 | 2.666e-06 | core.py:1476(strip_weak_type) |
| 76 | 0.0002596 | 3.416e-06 | 0.005257 | 6.917e-05 | util.py:298(promote_dtypes_numeric) |
| 47 | 0.0002586 | 5.502e-06 | 0.001298 | 2.763e-05 | _stablehlo_ops_gen.py:3549(__init__) |
| 154 | 0.0002574 | 1.671e-06 | 0.0002574 | 1.671e-06 | _func_ops_ext.py:71(body) |
| 705 | 0.000256 | 3.631e-07 | 0.0003686 | 5.229e-07 | inspect.py:755(_is_wrapper) |
| 335 | 0.000254 | 7.581e-07 | 0.000254 | 7.581e-07 | core.py:186(<listcomp>) |
| 500 | 0.0002533 | 5.066e-07 | 0.0003559 | 7.119e-07 | core.py:852(__hash__) |
| 1033 | 0.0002513 | 2.432e-07 | 0.0002513 | 2.432e-07 | effects.py:60(<listcomp>) |
| 2045 | 0.0002491 | 1.218e-07 | 0.0002491 | 1.218e-07 | partial_eval.py:127(get_known) |
| 42 | 0.0002476 | 5.896e-06 | 0.000781 | 1.859e-05 | lax.py:2897(_concatenate_shape_rule) |
| 25 | 0.0002476 | 9.906e-06 | 0.03907 | 0.001563 | reductions.py:219(sum) |
| 1110/658 | 0.0002475 | 3.762e-07 | 0.0007882 | 1.198e-06 | core.py:2236(_param_uses_outfeed) |
| 23 | 0.0002466 | 1.072e-05 | 0.01497 | 0.0006509 | partial_eval.py:2111(trace_to_jaxpr_final) |
| 37 | 0.0002465 | 6.663e-06 | 0.003414 | 9.228e-05 | slicing.py:737(_slice_shape_rule) |
| 154 | 0.0002461 | 1.598e-06 | 0.001324 | 8.599e-06 | util.py:150(unflatten) |
| 370 | 0.0002461 | 6.652e-07 | 0.0002461 | 6.652e-07 | __init__.py:1734(isEnabledFor) |
| 564 | 0.0002441 | 4.328e-07 | 0.001955 | 3.466e-06 | lax.py:1494(<genexpr>) |
| 664 | 0.000242 | 3.645e-07 | 0.0003083 | 4.642e-07 | data_structures.py:46(__len__) |
| 77 | 0.0002414 | 3.136e-06 | 0.0002414 | 3.136e-06 | _func_ops_ext.py:124(result_attrs) |
| 44 | 0.0002412 | 5.481e-06 | 0.0005519 | 1.254e-05 | pjit.py:2035(_fast_path_get_device_assignment) |
| 300 | 0.0002405 | 8.017e-07 | 0.0002405 | 8.017e-07 | ~:0(<method 'getbuffer' of '_io.BytesIO' objects>) |
| 159 | 0.0002395 | 1.507e-06 | 0.0092 | 5.786e-05 | partial_eval.py:2008(debug_info_final) |
| 1881 | 0.0002381 | 1.266e-07 | 0.0002381 | 1.266e-07 | partial_eval.py:1638(<genexpr>) |
| 136 | 0.000238 | 1.75e-06 | 0.000238 | 1.75e-06 | ~:0(<built-in method numpy.asanyarray>) |
| 1300 | 0.0002378 | 1.829e-07 | 0.0002378 | 1.829e-07 | core.py:390(get_bind_params) |
| 33 | 0.0002374 | 7.195e-06 | 0.007684 | 0.0002329 | lax.py:952(reduce) |
| 12 | 0.0002365 | 1.971e-05 | 0.0005289 | 4.408e-05 | _stablehlo_ops_gen.py:5053(__init__) |
| 1099 | 0.0002365 | 2.152e-07 | 0.0002365 | 2.152e-07 | linear_util.py:89(__init__) |
| 8 | 0.0002358 | 2.948e-05 | 0.06428 | 0.008035 | mlp.py:87(__call__) |
| 989 | 0.000235 | 2.376e-07 | 0.0003101 | 3.136e-07 | core.py:3012(<genexpr>) |
| 15 | 0.0002349 | 1.566e-05 | 0.0002349 | 1.566e-05 | ~:0(<built-in method importlib._bootstrap.get>) |
| 192 | 0.0002344 | 1.221e-06 | 0.000529 | 2.755e-06 | module.py:777(_module_method_call) |
| 62 | 0.000233 | 3.758e-06 | 0.000362 | 5.838e-06 | <frozen os>:674(__getitem__) |
| 170 | 0.0002301 | 1.354e-06 | 0.0009637 | 5.669e-06 | xla_bridge.py:649(process_count) |
| 78 | 0.0002289 | 2.934e-06 | 0.0002847 | 3.65e-06 | ~:0(<method 'sub' of 're.Pattern' objects>) |
| 250 | 0.0002286 | 9.143e-07 | 0.0002753 | 1.101e-06 | partial_eval.py:1582(<listcomp>) |
| 133 | 0.0002283 | 1.717e-06 | 0.00663 | 4.985e-05 | util.py:226(_asarray) |
| 1020 | 0.0002258 | 2.214e-07 | 0.0003808 | 3.733e-07 | pjit.py:1386(<genexpr>) |
| 318 | 0.0002232 | 7.018e-07 | 0.0002232 | 7.018e-07 | api_util.py:130(<genexpr>) |
| 576 | 0.0002227 | 3.866e-07 | 0.0004607 | 7.998e-07 | data_structures.py:76(__call__) |
| 77 | 0.0002209 | 2.869e-06 | 0.0002553 | 3.316e-06 | mlir.py:949(<listcomp>) |
| 558 | 0.000219 | 3.925e-07 | 0.001756 | 3.148e-06 | core.py:1553(strip_named_shape) |
| 598 | 0.0002187 | 3.657e-07 | 0.0003304 | 5.525e-07 | mlir.py:216(<genexpr>) |
| 560 | 0.0002179 | 3.891e-07 | 0.0002552 | 4.557e-07 | dispatch.py:417(eqn_replicas) |
| 108 | 0.000216 | 2e-06 | 0.001744 | 1.615e-05 | pxla.py:2330(_get_default_device) |
| 2 | 0.0002141 | 0.000107 | 0.0002481 | 0.000124 | pxla.py:262(op_sharding_to_indices) |
| 52 | 0.0002138 | 4.111e-06 | 0.01193 | 0.0002294 | ufuncs.py:329(power) |
| 57 | 0.0002119 | 3.717e-06 | 0.002223 | 3.9e-05 | lax.py:2594(precision_attr) |
| 487 | 0.0002097 | 4.306e-07 | 0.008115 | 1.666e-05 | ~:0(<method 'send' of 'generator' objects>) |
| 73 | 0.0002095 | 2.871e-06 | 0.0002095 | 2.871e-06 | _stablehlo_ops_gen.py:1247(result) |
| 4 | 0.0002089 | 5.222e-05 | 0.0006985 | 0.0001746 | util.py:208(check_toposort) |
| 794 | 0.0002081 | 2.621e-07 | 0.0002081 | 2.621e-07 | core.py:176(<lambda>) |
| 250 | 0.0002071 | 8.285e-07 | 0.0002071 | 8.285e-07 | partial_eval.py:2065(<listcomp>) |
| 60 | 0.0002068 | 3.447e-06 | 0.02315 | 0.0003858 | lax.py:1976(_integer_pow_lowering) |
| 76 | 0.0002044 | 2.69e-06 | 0.001407 | 1.852e-05 | _dtype.py:34(__str__) |
| 149 | 0.0002038 | 1.368e-06 | 0.0003678 | 2.468e-06 | core.py:1939(symbolic_equal_one_of_dim) |
| 280 | 0.0002036 | 7.273e-07 | 0.0002907 | 1.038e-06 | core.py:200(effects) |
| 300 | 0.0002031 | 6.769e-07 | 0.0002031 | 6.769e-07 | pickle.py:205(__init__) |
| 62 | 0.0002012 | 3.246e-06 | 0.001488 | 2.401e-05 | mlir.py:684(<listcomp>) |
| 1154 | 0.0001999 | 1.732e-07 | 0.0002872 | 2.488e-07 | ~:0(<method 'extend' of 'list' objects>) |
| 300 | 0.0001998 | 6.661e-07 | 0.0003217 | 1.072e-06 | <frozen posixpath>:41(_get_sep) |
| 62 | 0.0001991 | 3.211e-06 | 0.0001991 | 3.211e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.append_dialect_registry>) |
| 62 | 0.0001986 | 3.203e-06 | 4.125 | 0.06653 | dispatch.py:461(backend_compile) |
| 64 | 0.0001986 | 3.103e-06 | 0.00192 | 3e-05 | pxla.py:1731(global_avals_to_results_handler) |
| 88 | 0.0001981 | 2.251e-06 | 0.0004446 | 5.053e-06 | pjit.py:1289(__hash__) |
| 10 | 0.000198 | 1.98e-05 | 0.0004688 | 4.688e-05 | _stablehlo_ops_gen.py:4717(__init__) |
| 269 | 0.000197 | 7.322e-07 | 0.0006019 | 2.237e-06 | util.py:426(__call__) |
| 545 | 0.0001969 | 3.613e-07 | 0.00168 | 3.083e-06 | pjit.py:467(<lambda>) |
| 126 | 0.0001967 | 1.561e-06 | 0.001003 | 7.959e-06 | sharding.py:82(is_fully_addressable) |
| 199 | 0.0001965 | 9.872e-07 | 0.00202 | 1.015e-05 | util.py:357(<genexpr>) |
| 758 | 0.0001955 | 2.579e-07 | 0.0001955 | 2.579e-07 | core.py:844(__init__) |
| 62 | 0.0001954 | 3.152e-06 | 0.000423 | 6.823e-06 | mlir.py:818(create) |
| 250 | 0.0001951 | 7.802e-07 | 0.0003858 | 1.543e-06 | partial_eval.py:1709(<listcomp>) |
| 56/41 | 0.0001946 | 4.747e-06 | 0.001385 | 3.379e-05 | dispatch.py:309(jaxpr_has_primitive) |
| 179 | 0.0001932 | 1.079e-06 | 0.0002986 | 1.668e-06 | lax.py:2720(_broadcast_in_dim_fwd_rule) |
| 24 | 0.0001925 | 8.019e-06 | 0.06622 | 0.002759 | transform.py:348(<lambda>) |
| 132 | 0.0001925 | 1.458e-06 | 0.0008279 | 6.272e-06 | partial_eval.py:155(new_const) |
| 77 | 0.0001915 | 2.487e-06 | 0.0001915 | 2.487e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.insert>) |
| 642 | 0.0001912 | 2.978e-07 | 0.0001912 | 2.978e-07 | ~:0(<method 'values' of 'mappingproxy' objects>) |
| 62 | 0.0001912 | 3.083e-06 | 0.0003624 | 5.846e-06 | pxla.py:2942(load) |
| 99 | 0.0001893 | 1.912e-06 | 0.0007832 | 7.911e-06 | partial_eval.py:165(new_instantiated_const) |
| 15 | 0.0001886 | 1.257e-05 | 0.00362 | 0.0002413 | pjit.py:1769(dce_jaxpr_pjit_rule) |
| 94/38 | 0.0001876 | 4.937e-06 | 0.0002986 | 7.858e-06 | pprint.py:551(_safe_repr) |
| 545 | 0.0001864 | 3.42e-07 | 0.0001864 | 3.42e-07 | api_util.py:617(<dictcomp>) |
| 48 | 0.0001859 | 3.872e-06 | 0.0008809 | 1.835e-05 | module.py:457(unique_and_canonical_name) |
| 340 | 0.0001859 | 5.468e-07 | 0.0001859 | 5.468e-07 | xla_bridge.py:653(<genexpr>) |
| 186 | 0.0001853 | 9.961e-07 | 0.000283 | 1.521e-06 | functools.py:89(_gt_from_lt) |
| 77 | 0.0001838 | 2.387e-06 | 0.000978 | 1.27e-05 | mlir.py:995(<listcomp>) |
| 186 | 0.0001833 | 9.856e-07 | 0.0002279 | 1.225e-06 | effects.py:62(filter_not_in) |
| 57 | 0.0001804 | 3.166e-06 | 0.00148 | 2.596e-05 | lax.py:2602(<listcomp>) |
| 73 | 0.0001783 | 2.443e-06 | 0.00362 | 4.959e-05 | lax.py:2283(_convert_element_type_lower) |
| 951 | 0.0001761 | 1.852e-07 | 0.0001761 | 1.852e-07 | utils.py:118(<genexpr>) |
| 168 | 0.0001758 | 1.046e-06 | 0.00746 | 4.44e-05 | lax.py:2446(_dot_general_dtype_rule) |
| 62 | 0.0001756 | 2.832e-06 | 0.000252 | 4.065e-06 | pxla.py:1855(__init__) |
| 971 | 0.0001744 | 1.796e-07 | 0.0001744 | 1.796e-07 | mlir.py:832(<genexpr>) |
| 64 | 0.000174 | 2.719e-06 | 0.000174 | 2.719e-06 | pxla.py:3121(__init__) |
| 62 | 0.000172 | 2.775e-06 | 0.0005416 | 8.735e-06 | function_base.py:2331(_get_ufunc_and_otypes) |
| 668 | 0.0001713 | 2.564e-07 | 0.0001713 | 2.564e-07 | lax.py:2443(<genexpr>) |
| 278 | 0.0001709 | 6.146e-07 | 0.0005361 | 1.928e-06 | xla_bridge.py:559(devices) |
| 64 | 0.0001698 | 2.654e-06 | 0.0002044 | 3.194e-06 | pxla.py:2456(<listcomp>) |
| 64 | 0.0001691 | 2.642e-06 | 0.0001691 | 2.642e-06 | pxla.py:2420(<listcomp>) |
| 281 | 0.000167 | 5.944e-07 | 0.0004628 | 1.647e-06 | lax.py:2656(<genexpr>) |
| 244 | 0.0001666 | 6.829e-07 | 0.0001666 | 6.829e-07 | monitoring.py:33(record_event_duration_secs) |
| 149 | 0.0001662 | 1.115e-06 | 0.009496 | 6.373e-05 | lax.py:484(convert_element_type) |
| 728 | 0.0001654 | 2.272e-07 | 0.0001654 | 2.272e-07 | data_structures.py:73(peek) |
| 523 | 0.0001653 | 3.161e-07 | 0.0005472 | 1.046e-06 | pjit.py:858(<genexpr>) |
| 43 | 0.0001649 | 3.834e-06 | 0.01366 | 0.0003176 | util.py:368(promote_args) |
| 315 | 0.0001641 | 5.21e-07 | 0.0007194 | 2.284e-06 | partial_eval.py:120(unknown) |
| 8 | 0.0001639 | 2.049e-05 | 0.05045 | 0.006306 | reductions.py:427(_var) |
| 62 | 0.0001639 | 2.644e-06 | 0.0003525 | 5.685e-06 | xla_bridge.py:365(canonicalize_platform) |
| 26 | 0.0001636 | 6.294e-06 | 0.002335 | 8.982e-05 | lax.py:2971(_concatenate_lower) |
| 1277 | 0.0001633 | 1.279e-07 | 0.0001633 | 1.279e-07 | core.py:92(<lambda>) |
| 112 | 0.0001632 | 1.457e-06 | 0.0001883 | 1.681e-06 | __init__.py:585(__init__) |
| 299 | 0.0001626 | 5.437e-07 | 0.0001626 | 5.437e-07 | ~:0(<built-in method numpy.ascontiguousarray>) |
| 250 | 0.0001624 | 6.494e-07 | 0.0002975 | 1.19e-06 | partial_eval.py:1679(<listcomp>) |
| 107 | 0.0001605 | 1.5e-06 | 0.0002002 | 1.871e-06 | lax.py:2248(_convert_elt_type_fwd_rule) |
| 114/2 | 0.0001603 | 8.013e-05 | 0.06521 | 0.03261 | partial_eval.py:209(process_primitive) |
| 423 | 0.0001601 | 3.784e-07 | 0.0001601 | 3.784e-07 | ~:0(<built-in method jaxlib.xla_extension.is_replicated>) |
| 558 | 0.0001595 | 2.858e-07 | 0.0002789 | 4.998e-07 | pxla.py:2432(<genexpr>) |
| 48 | 0.0001593 | 3.318e-06 | 0.001369 | 2.852e-05 | module.py:597(__init__) |
| 64 | 0.000158 | 2.469e-06 | 0.000158 | 2.469e-06 | pxla.py:1667(__init__) |
| 731 | 0.0001559 | 2.133e-07 | 0.0002312 | 3.163e-07 | pjit.py:1290(<genexpr>) |
| 551 | 0.0001558 | 2.828e-07 | 0.0001558 | 2.828e-07 | core.py:2588(<genexpr>) |
| 159 | 0.0001555 | 9.78e-07 | 0.0006324 | 3.978e-06 | api_util.py:629(<listcomp>) |
| 1116 | 0.0001553 | 1.391e-07 | 0.0001553 | 1.391e-07 | core.py:91(<lambda>) |
| 96/16 | 0.0001552 | 9.699e-06 | 0.08407 | 0.005254 | module.py:263(run_interceptors) |
| 545 | 0.0001552 | 2.848e-07 | 0.0001552 | 2.848e-07 | inspect.py:2828(__init__) |
| 47 | 0.0001548 | 3.295e-06 | 0.0008973 | 1.909e-05 | _stablehlo_ops_gen.py:3952(__init__) |
| 17 | 0.0001505 | 8.852e-06 | 0.034 | 0.002 | reductions.py:326(_mean) |
| 2 | 0.0001486 | 7.428e-05 | 0.06519 | 0.0326 | pjit.py:1572(_pjit_partial_eval) |
| 346 | 0.000147 | 4.249e-07 | 0.000147 | 4.249e-07 | core.py:190(<listcomp>) |
| 76 | 0.0001458 | 1.919e-06 | 0.0009322 | 1.227e-05 | fromnumeric.py:38(_wrapit) |
| 192 | 0.0001444 | 7.518e-07 | 0.0001706 | 8.887e-07 | data_structures.py:49(__iter__) |
| 48 | 0.000144 | 3e-06 | 0.0003269 | 6.81e-06 | module.py:493(reserve_module_name) |
| 30 | 0.0001438 | 4.793e-06 | 0.008318 | 0.0002773 | lax.py:693(dot_general) |
| 23 | 0.0001431 | 6.221e-06 | 1.059 | 0.04602 | dispatch.py:249(_xla_callable_uncached) |
| 441 | 0.0001429 | 3.24e-07 | 0.0001871 | 4.242e-07 | lax.py:2808(<genexpr>) |
| 64 | 0.0001416 | 2.212e-06 | 0.0005245 | 8.195e-06 | source_info_util.py:106(new_name_stack) |
| 62 | 0.0001403 | 2.263e-06 | 0.0001403 | 2.263e-06 | pxla.py:2876(<listcomp>) |
| 49/9 | 0.0001387 | 1.541e-05 | 0.07368 | 0.008187 | contextlib.py:78(inner) |
| 112 | 0.0001382 | 1.234e-06 | 0.0004412 | 3.939e-06 | partial_eval.py:465(_current_truncated_name_stack) |
| 209 | 0.0001371 | 6.561e-07 | 0.0001984 | 9.491e-07 | ad.py:207(read_cotangent) |
| 114 | 0.0001359 | 1.193e-06 | 0.0002857 | 2.506e-06 | partial_eval.py:862(<listcomp>) |
| 64 | 0.0001354 | 2.116e-06 | 0.0001354 | 2.116e-06 | ~:0(<method 'reshape' of 'numpy.ndarray' objects>) |
| 205 | 0.0001342 | 6.544e-07 | 0.0008217 | 4.008e-06 | partial_eval.py:191(instantiate_const) |
| 62 | 0.000134 | 2.162e-06 | 0.0001754 | 2.829e-06 | mlir.py:1201(_ir_consts) |
| 141 | 0.0001332 | 9.444e-07 | 0.001319 | 9.357e-06 | pxla.py:586(global_aval_to_result_handler) |
| 159 | 0.0001312 | 8.249e-07 | 0.09743 | 0.0006128 | core.py:645(__add__) |
| 235 | 0.0001305 | 5.551e-07 | 0.00017 | 7.234e-07 | core.py:880(pop) |
| 84 | 0.0001302 | 1.55e-06 | 0.0007473 | 8.897e-06 | lax.py:1605(_unbroadcast) |
| 124 | 0.0001299 | 1.048e-06 | 0.0001842 | 1.485e-06 | pxla.py:3042(<genexpr>) |
| 62 | 0.0001299 | 2.096e-06 | 0.0001299 | 2.096e-06 | function_base.py:2286(<listcomp>) |
| 66 | 0.0001295 | 1.963e-06 | 0.0009308 | 1.41e-05 | lax.py:1942(_integer_pow_dtype_rule) |
| 62 | 0.0001294 | 2.087e-06 | 0.0001486 | 2.397e-06 | xla_bridge.py:397(backends) |
| 328 | 0.0001292 | 3.94e-07 | 0.0001526 | 4.651e-07 | partial_eval.py:898(get_atom) |
| 62 | 0.000129 | 2.08e-06 | 0.001636 | 2.639e-05 | function_base.py:2301(__call__) |
| 360 | 0.0001277 | 3.547e-07 | 0.0001584 | 4.399e-07 | data_structures.py:64(push) |
| 44 | 0.0001261 | 2.866e-06 | 0.0001261 | 2.866e-06 | pjit.py:1161(<listcomp>) |
| 114 | 0.0001249 | 1.096e-06 | 0.0003277 | 2.874e-06 | lax.py:645(__init__) |
| 128 | 0.0001248 | 9.75e-07 | 0.0002232 | 1.744e-06 | base.py:287(current_module_state) |
| 250 | 0.0001248 | 4.991e-07 | 0.0001593 | 6.373e-07 | partial_eval.py:1659(<listcomp>) |
| 300 | 0.0001247 | 4.156e-07 | 0.0001247 | 4.156e-07 | ~:0(<built-in method posix.fspath>) |
| 85 | 0.0001245 | 1.464e-06 | 0.0008356 | 9.831e-06 | lax.py:3508(_reduce_number_dtype_rule) |
| 48 | 0.0001244 | 2.591e-06 | 0.000222 | 4.624e-06 | utils.py:79(<listcomp>) |
| 250 | 0.0001235 | 4.939e-07 | 0.0001448 | 5.793e-07 | partial_eval.py:1701(<listcomp>) |
| 62 | 0.0001235 | 1.992e-06 | 0.0005751 | 9.276e-06 | dispatch.py:287(raise_warnings_or_errors_for_jit_of_pmap) |
| 62 | 0.0001233 | 1.989e-06 | 0.0004853 | 7.827e-06 | <frozen _collections_abc>:771(get) |
| 442 | 0.000123 | 2.784e-07 | 0.000123 | 2.784e-07 | util.py:300(<genexpr>) |
| 194 | 0.000123 | 6.342e-07 | 0.0003023 | 1.558e-06 | config.py:541(__getattr__) |
| 194 | 0.0001224 | 6.307e-07 | 0.0001792 | 9.239e-07 | config.py:87(read) |
| 144 | 0.0001211 | 8.408e-07 | 0.0004993 | 3.467e-06 | source_info_util.py:234(extend_name_stack) |
| 108 | 0.0001206 | 1.116e-06 | 0.0001206 | 1.116e-06 | xla_bridge.py:521(get_device_backend) |
| 42 | 0.0001205 | 2.87e-06 | 0.001049 | 2.497e-05 | lax.py:2925(_concatenate_dtype_rule) |
| 742 | 0.0001196 | 1.611e-07 | 0.0001196 | 1.611e-07 | core.py:1820(symbolic_equal) |
| 237 | 0.0001194 | 5.037e-07 | 0.0006853 | 2.892e-06 | api_util.py:88(flatten_fun) |
| 243 | 0.0001193 | 4.908e-07 | 0.0117 | 4.814e-05 | pxla.py:3137(unsafe_call) |
| 46 | 0.0001175 | 2.555e-06 | 0.0008568 | 1.863e-05 | core.py:1110(new_base_main) |
| 64 | 0.0001168 | 1.826e-06 | 0.001436 | 2.244e-05 | pxla.py:1736(<listcomp>) |
| 12 | 0.0001167 | 9.725e-06 | 0.003523 | 0.0002936 | lax.py:2687(_broadcast_in_dim_transpose_rule) |
| 118 | 0.0001164 | 9.861e-07 | 0.0001643 | 1.392e-06 | partial_eval.py:142(__init__) |
| 6 | 0.0001155 | 1.925e-05 | 0.0002595 | 4.324e-05 | _stablehlo_ops_gen.py:2963(__init__) |
| 132 | 0.0001145 | 8.676e-07 | 0.0001145 | 8.676e-07 | getlimits.py:679(min) |
| 37 | 0.000114 | 3.081e-06 | 0.001378 | 3.724e-05 | lax.py:1292(squeeze) |
| 62 | 0.0001139 | 1.836e-06 | 0.000136 | 2.194e-06 | mlir.py:963(<listcomp>) |
| 25 | 0.0001112 | 4.447e-06 | 0.001219 | 4.875e-05 | lax.py:929(transpose) |
| 224 | 0.0001112 | 4.963e-07 | 0.0004676 | 2.088e-06 | partial_eval.py:116(known) |
| 8 | 0.0001109 | 1.386e-05 | 0.007859 | 0.0009824 | mlp.py:41(__init__) |
| 12 | 0.0001096 | 9.137e-06 | 0.0006258 | 5.215e-05 | mlir.py:1442(_minmax_hlo) |
| 6 | 0.0001092 | 1.819e-05 | 0.0002493 | 4.156e-05 | _stablehlo_ops_gen.py:3000(__init__) |
| 112/16 | 0.000109 | 6.813e-06 | 0.0007422 | 4.639e-05 | data_structures.py:148(_to_dict_recurse) |
| 235 | 0.000109 | 4.637e-07 | 0.0001424 | 6.061e-07 | core.py:874(next_level) |
| 62 | 0.0001079 | 1.74e-06 | 0.0001623 | 2.617e-06 | xla_bridge.py:530(device_count) |
| 96 | 0.0001077 | 1.122e-06 | 0.0001431 | 1.49e-06 | contextlib.py:546(__exit__) |
| 76 | 0.0001044 | 1.373e-06 | 0.0001044 | 1.373e-06 | dtypes.py:81(to_numeric_dtype) |
| 255 | 0.0001042 | 4.084e-07 | 0.0006647 | 2.607e-06 | api_util.py:102(flatten_fun_nokwargs) |
| 235 | 0.0001035 | 4.404e-07 | 0.0001294 | 5.505e-07 | core.py:877(push) |
| 41 | 0.0001032 | 2.518e-06 | 0.000193 | 4.707e-06 | pxla.py:2261(check_if_any_auto) |
| 62 | 0.0001032 | 1.664e-06 | 0.0005797 | 9.35e-06 | dispatch.py:405(jaxpr_replicas) |
| 44 | 0.0001029 | 2.338e-06 | 0.0002034 | 4.622e-06 | pxla.py:3323(check_device_backend_on_shardings) |
| 24 | 0.0001026 | 4.276e-06 | 0.04507 | 0.001878 | transform.py:107(<lambda>) |
| 96 | 0.0001025 | 1.067e-06 | 0.0003537 | 3.684e-06 | module.py:449(name_and_number) |
| 485 | 0.0001024 | 2.111e-07 | 0.0001024 | 2.111e-07 | pxla.py:2451(<genexpr>) |
| 64 | 0.0001019 | 1.593e-06 | 0.001403 | 2.193e-05 | base.py:446(wrapped) |
| 33 | 0.0001011 | 3.064e-06 | 0.01579 | 0.0004784 | ufuncs.py:94(fn) |
| 114 | 0.000101 | 8.857e-07 | 0.0004287 | 3.76e-06 | lax.py:607(__get__) |
| 188 | 9.967e-05 | 5.301e-07 | 0.0002614 | 1.39e-06 | ad.py:453(__init__) |
| 8 | 9.947e-05 | 1.243e-05 | 0.0002153 | 2.691e-05 | pretty_printer.py:275(_format) |
| 6 | 9.931e-05 | 1.655e-05 | 0.0002354 | 3.923e-05 | _stablehlo_ops_gen.py:3388(__init__) |
| 165 | 9.837e-05 | 5.962e-07 | 0.0001265 | 7.667e-07 | core.py:2553(<listcomp>) |
| 46 | 9.834e-05 | 2.138e-06 | 0.0001373 | 2.985e-06 | core.py:1197(new_sublevel) |
| 62 | 9.832e-05 | 1.586e-06 | 9.832e-05 | 1.586e-06 | xla_bridge.py:143(<lambda>) |
| 312 | 9.829e-05 | 3.15e-07 | 0.0001243 | 3.983e-07 | data_structures.py:70(pop) |
| 141 | 9.784e-05 | 6.939e-07 | 9.784e-05 | 6.939e-07 | util.py:324(wrap_name) |
| 2 | 9.77e-05 | 4.885e-05 | 0.0599 | 0.02995 | partial_eval.py:1061(_partial_eval_jaxpr_nounits) |
| 60 | 9.663e-05 | 1.61e-06 | 0.005237 | 8.728e-05 | lax.py:1958(_integer_pow) |
| 124 | 9.628e-05 | 7.764e-07 | 0.08413 | 0.0006785 | core.py:650(__rmul__) |
| 24 | 9.6e-05 | 4e-06 | 0.01627 | 0.0006781 | update.py:43(<lambda>) |
| 76 | 9.567e-05 | 1.259e-06 | 0.001038 | 1.366e-05 | fromnumeric.py:51(_wrapfunc) |
| 23 | 9.541e-05 | 4.148e-06 | 0.06154 | 0.002675 | dispatch.py:235(sharded_lowering) |
| 25 | 9.528e-05 | 3.811e-06 | 0.01303 | 0.000521 | reductions.py:208(_reduce_sum) |
| 48 | 9.505e-05 | 1.98e-06 | 0.0003598 | 7.496e-06 | module.py:385(wrap_method) |
| 112 | 9.377e-05 | 8.372e-07 | 0.0001224 | 1.092e-06 | partial_eval.py:219(<listcomp>) |
| 26 | 9.337e-05 | 3.591e-06 | 0.005652 | 0.0002174 | util.py:381(promote_args_inexact) |
| 642 | 9.331e-05 | 1.453e-07 | 9.331e-05 | 1.453e-07 | partial_eval.py:1640(<genexpr>) |
| 47 | 9.255e-05 | 1.969e-06 | 9.255e-05 | 1.969e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.append>) |
| 627 | 9.235e-05 | 1.473e-07 | 9.235e-05 | 1.473e-07 | tree_util.py:533(__str__) |
| 188 | 9.175e-05 | 4.88e-07 | 0.0001426 | 7.584e-07 | lax.py:3553(<genexpr>) |
| 48 | 9.173e-05 | 1.911e-06 | 0.0003197 | 6.659e-06 | partial_eval.py:174(new_arg) |
| 321 | 9.134e-05 | 2.846e-07 | 9.134e-05 | 2.846e-07 | lax.py:2198(_convert_element_type_shape_rule) |
| 250 | 9.116e-05 | 3.646e-07 | 0.0001163 | 4.653e-07 | partial_eval.py:1700(<listcomp>) |
| 343 | 9.116e-05 | 2.658e-07 | 0.0001247 | 3.635e-07 | tree_util.py:521(__str__) |
| 554 | 8.999e-05 | 1.624e-07 | 8.999e-05 | 1.624e-07 | dtypes.py:661(<genexpr>) |
| 3 | 8.979e-05 | 2.993e-05 | 0.2166 | 0.0722 | transform.py:340(update_fn) |
| 58 | 8.893e-05 | 1.533e-06 | 0.0007708 | 1.329e-05 | ad.py:302(lift) |
| 539 | 8.859e-05 | 1.644e-07 | 0.0001076 | 1.996e-07 | util.py:211(<genexpr>) |
| 568 | 8.855e-05 | 1.559e-07 | 8.855e-05 | 1.559e-07 | ~:0(<method 'pop' of 'dict' objects>) |
| 17 | 8.837e-05 | 5.198e-06 | 0.0004613 | 2.713e-05 | pjit.py:889(_check_and_canonicalize_out_shardings) |
| 99 | 8.788e-05 | 8.877e-07 | 0.0007338 | 7.412e-06 | ~:0(<built-in method _functools.reduce>) |
| 79/64 | 8.706e-05 | 1.36e-06 | 0.001636 | 2.557e-05 | core.py:2231(jaxpr_uses_outfeed) |
| 195 | 8.696e-05 | 4.459e-07 | 0.0003119 | 1.6e-06 | linear_util.py:211(transformation) |
| 45 | 8.635e-05 | 1.919e-06 | 0.0002073 | 4.606e-06 | pjit.py:1774(keep_where) |
| 44 | 8.633e-05 | 1.962e-06 | 8.633e-05 | 1.962e-06 | pjit.py:257(<listcomp>) |
| 8 | 8.62e-05 | 1.078e-05 | 9.835e-05 | 1.229e-05 | util.py:320(_complex_elem_type) |
| 2 | 8.602e-05 | 4.301e-05 | 0.0003061 | 0.000153 | partial_eval.py:1323(_jaxpr_forwarding) |
| 42 | 8.577e-05 | 2.042e-06 | 0.0001775 | 4.225e-06 | lax.py:2906(<setcomp>) |
| 86 | 8.546e-05 | 9.937e-07 | 0.0007586 | 8.821e-06 | core.py:1530(at_least_vspace) |
| 10 | 8.486e-05 | 8.486e-06 | 0.0008119 | 8.119e-05 | mlir.py:1342(slice_op) |
| 24 | 8.374e-05 | 3.489e-06 | 0.0573 | 0.002387 | transform.py:84(<lambda>) |
| 1048 | 8.353e-05 | 7.97e-08 | 8.353e-05 | 7.97e-08 | inspect.py:2744(annotation) |
| 23 | 8.275e-05 | 3.598e-06 | 1.064 | 0.04627 | util.py:244(cached) |
| 545 | 8.27e-05 | 1.517e-07 | 8.27e-05 | 1.517e-07 | pjit.py:544(_flat_axes_specs) |
| 33 | 8.258e-05 | 2.502e-06 | 0.0003381 | 1.025e-05 | reductions.py:150(_reduction_dims) |
| 77 | 8.112e-05 | 1.054e-06 | 8.112e-05 | 1.054e-06 | partial_eval.py:1388(<listcomp>) |
| 85 | 8.083e-05 | 9.51e-07 | 0.000646 | 7.6e-06 | lax.py:3514(_reduce_sum_shape_rule) |
| 120 | 8.067e-05 | 6.723e-07 | 0.0001282 | 1.068e-06 | ~:0(<built-in method builtins.sum>) |
| 645 | 8.046e-05 | 1.247e-07 | 8.046e-05 | 1.247e-07 | inspect.py:3031(parameters) |
| 41 | 8.03e-05 | 1.959e-06 | 0.0001632 | 3.98e-06 | typing.py:345(inner) |
| 485 | 8.006e-05 | 1.651e-07 | 8.006e-05 | 1.651e-07 | pxla.py:2452(<genexpr>) |
| 101 | 8.003e-05 | 7.923e-07 | 8.003e-05 | 7.923e-07 | api_util.py:330(<dictcomp>) |
| 423 | 7.993e-05 | 1.89e-07 | 7.993e-05 | 1.89e-07 | core.py:1940(<genexpr>) |
| 22 | 7.984e-05 | 3.629e-06 | 0.01204 | 0.0005475 | ad.py:524(standard_jvp) |
| 269 | 7.983e-05 | 2.968e-07 | 0.0001095 | 4.072e-07 | util.py:287(__hash__) |
| 250 | 7.923e-05 | 3.169e-07 | 7.923e-05 | 3.169e-07 | partial_eval.py:2071(<listcomp>) |
| 42 | 7.917e-05 | 1.885e-06 | 7.917e-05 | 1.885e-06 | lax.py:2912(<listcomp>) |
| 336 | 7.902e-05 | 2.352e-07 | 7.902e-05 | 2.352e-07 | lax.py:2423(<genexpr>) |
| 128 | 7.864e-05 | 6.144e-07 | 0.0003185 | 2.488e-06 | typing.py:1566(__subclasscheck__) |
| 242 | 7.815e-05 | 3.229e-07 | 8.904e-05 | 3.679e-07 | pjit.py:251(<genexpr>) |
| 112 | 7.804e-05 | 6.968e-07 | 0.0002285 | 2.04e-06 | partial_eval.py:223(<listcomp>) |
| 112 | 7.791e-05 | 6.956e-07 | 0.0001406 | 1.255e-06 | source_info_util.py:87(__getitem__) |
| 107 | 7.696e-05 | 7.193e-07 | 0.01524 | 0.0001424 | lax.py:408(mul) |
| 62 | 7.693e-05 | 1.241e-06 | 0.000129 | 2.081e-06 | <frozen os>:756(encode) |
| 9 | 7.682e-05 | 8.535e-06 | 0.0003091 | 3.435e-05 | ad.py:86(jvp_subtrace) |
| 214 | 7.668e-05 | 3.583e-07 | 9.536e-05 | 4.456e-07 | partial_eval.py:1642(<listcomp>) |
| 108 | 7.667e-05 | 7.099e-07 | 7.667e-05 | 7.099e-07 | xla_bridge.py:619(<listcomp>) |
| 15 | 7.663e-05 | 5.108e-06 | 7.663e-05 | 5.108e-06 | _func_ops_ext.py:83(name) |
| 485 | 7.625e-05 | 1.572e-07 | 7.625e-05 | 1.572e-07 | pxla.py:2453(<genexpr>) |
| 8 | 7.561e-05 | 9.451e-06 | 0.05969 | 0.007461 | reductions.py:420(var) |
| 59 | 7.54e-05 | 1.278e-06 | 0.0005888 | 9.98e-06 | util.py:240(<listcomp>) |
| 8 | 7.496e-05 | 9.369e-06 | 0.0007002 | 8.753e-05 | base.py:806(assert_is_prng_key) |
| 48 | 7.484e-05 | 1.559e-06 | 0.0003864 | 8.049e-06 | contextlib.py:77(__call__) |
| 76 | 7.471e-05 | 9.83e-07 | 0.002735 | 3.598e-05 | util.py:305(<listcomp>) |
| 930 | 7.469e-05 | 8.031e-08 | 7.469e-05 | 8.031e-08 | inspect.py:2740(default) |
| 104 | 7.465e-05 | 7.178e-07 | 0.0001322 | 1.271e-06 | base.py:269(safe_get_module_name) |
| 510 | 7.445e-05 | 1.46e-07 | 7.445e-05 | 1.46e-07 | partial_eval.py:1652(<genexpr>) |
| 12 | 7.425e-05 | 6.188e-06 | 0.0007069 | 5.891e-05 | lax.py:3234(_transpose_lower) |
| 441 | 7.392e-05 | 1.676e-07 | 7.392e-05 | 1.676e-07 | lax.py:2806(<genexpr>) |
| 75 | 7.38e-05 | 9.84e-07 | 8.241e-05 | 1.099e-06 | api_util.py:320(<listcomp>) |
| 574 | 7.37e-05 | 1.284e-07 | 7.37e-05 | 1.284e-07 | core.py:1826(greater_equal) |
| 38 | 7.365e-05 | 1.938e-06 | 0.0006428 | 1.691e-05 | pprint.py:57(pformat) |
| 38 | 7.328e-05 | 1.928e-06 | 0.0008148 | 2.144e-05 | utils.py:97(fancy_repr) |
| 33 | 7.295e-05 | 2.211e-06 | 0.0005561 | 1.685e-05 | reductions.py:165(_reduction_init_val) |
| 188 | 7.291e-05 | 3.878e-07 | 0.0001175 | 6.249e-07 | lax.py:1511(<lambda>) |
| 321 | 7.264e-05 | 2.263e-07 | 9.069e-05 | 2.825e-07 | lax.py:2902(<genexpr>) |
| 29 | 7.261e-05 | 2.504e-06 | 0.004391 | 0.0001514 | lax.py:2079(_mul_transpose) |
| 41 | 7.252e-05 | 1.769e-06 | 8.286e-05 | 2.021e-06 | typing.py:1656(__hash__) |
| 190 | 7.195e-05 | 3.787e-07 | 0.0001145 | 6.028e-07 | partial_eval.py:1760(_lift_tracers_in_aval) |
| 79 | 7.179e-05 | 9.087e-07 | 0.0001747 | 2.211e-06 | ad_util.py:79(from_value) |
| 46 | 7.173e-05 | 1.559e-06 | 0.03897 | 0.0008472 | lax.py:3067(_squeeze_shape_rule) |
| 77 | 7.169e-05 | 9.31e-07 | 8.317e-05 | 1.08e-06 | mlir.py:954(<listcomp>) |
| 90 | 7.143e-05 | 7.936e-07 | 0.0004791 | 5.323e-06 | core.py:1623(primal_dtype_to_tangent_dtype) |
| 21 | 7.067e-05 | 3.365e-06 | 0.0004469 | 2.128e-05 | sharding.py:89(addressable_devices_indices_map) |
| 96 | 7.058e-05 | 7.352e-07 | 7.058e-05 | 7.352e-07 | ~:0(<method 'split' of 're.Pattern' objects>) |
| 25 | 6.955e-05 | 2.782e-06 | 0.009638 | 0.0003855 | ad.py:557(bilinear_transpose) |
| 30 | 6.954e-05 | 2.318e-06 | 6.954e-05 | 2.318e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.get_signless>) |
| 77 | 6.944e-05 | 9.018e-07 | 0.0003194 | 4.147e-06 | mlir.py:1012(<listcomp>) |
| 96 | 6.943e-05 | 7.232e-07 | 0.0002408 | 2.509e-06 | __init__.py:198(split) |
| 44 | 6.928e-05 | 1.574e-06 | 6.928e-05 | 1.574e-06 | pjit.py:255(<listcomp>) |
| 15 | 6.927e-05 | 4.618e-06 | 0.0005364 | 3.576e-05 | _func_ops_gen.py:64(__init__) |
| 48 | 6.917e-05 | 1.441e-06 | 0.003138 | 6.538e-05 | array_methods.py:54(_astype) |
| 22 | 6.912e-05 | 3.142e-06 | 0.006898 | 0.0003136 | ad.py:526(<listcomp>) |
| 44 | 6.909e-05 | 1.57e-06 | 6.909e-05 | 1.57e-06 | pjit.py:256(<listcomp>) |
| 26 | 6.86e-05 | 2.639e-06 | 0.00234 | 9.001e-05 | util.py:287(promote_dtypes_inexact) |
| 15 | 6.855e-05 | 4.57e-06 | 0.003257 | 0.0002171 | lax.py:1318(full_like) |
| 28 | 6.839e-05 | 2.443e-06 | 0.00022 | 7.857e-06 | mlir.py:80(i64_attr) |
| 728 | 6.708e-05 | 9.214e-08 | 6.708e-05 | 9.214e-08 | partial_eval.py:124(is_known) |
| 274 | 6.689e-05 | 2.441e-07 | 9.236e-05 | 3.371e-07 | pjit.py:1300(<genexpr>) |
| 222 | 6.678e-05 | 3.008e-07 | 9.34e-05 | 4.207e-07 | api_util.py:159(__hash__) |
| 96 | 6.622e-05 | 6.898e-07 | 0.0001008 | 1.05e-06 | __init__.py:272(_compile) |
| 274 | 6.577e-05 | 2.4e-07 | 0.0001079 | 3.937e-07 | pxla.py:2433(<genexpr>) |
| 38 | 6.553e-05 | 1.724e-06 | 0.0004534 | 1.193e-05 | pprint.py:168(_format) |
| 150 | 6.531e-05 | 4.354e-07 | 0.0002222 | 1.481e-06 | ad.py:460(aval) |
| 160 | 6.503e-05 | 4.064e-07 | 6.503e-05 | 4.064e-07 | ~:0(<method 'issubset' of 'set' objects>) |
| 442 | 6.496e-05 | 1.47e-07 | 6.496e-05 | 1.47e-07 | ad_util.py:75(__init__) |
| 8 | 6.478e-05 | 8.098e-06 | 0.005409 | 0.0006761 | partial_eval.py:789(_trace_to_subjaxpr_nounits) |
| 62 | 6.455e-05 | 1.041e-06 | 0.0198 | 0.0003194 | mlir.py:1213(<lambda>) |
| 64 | 6.443e-05 | 1.007e-06 | 6.443e-05 | 1.007e-06 | pxla.py:2781(__init__) |
| 46 | 6.397e-05 | 1.391e-06 | 0.002447 | 5.32e-05 | lax.py:2786(_broadcast_in_dim_lower) |
| 50 | 6.334e-05 | 1.267e-06 | 8.795e-05 | 1.759e-06 | lax.py:4618(remaining) |
| 16 | 6.261e-05 | 3.913e-06 | 0.01566 | 0.0009789 | transform.py:122(<lambda>) |
| 64 | 6.241e-05 | 9.751e-07 | 6.241e-05 | 9.751e-07 | dispatch.py:374(<setcomp>) |
| 210 | 6.238e-05 | 2.97e-07 | 6.238e-05 | 2.97e-07 | api_util.py:294(<genexpr>) |
| 250 | 6.235e-05 | 2.494e-07 | 6.235e-05 | 2.494e-07 | partial_eval.py:1676(<setcomp>) |
| 128 | 6.194e-05 | 4.839e-07 | 0.0003804 | 2.972e-06 | typing.py:1291(__instancecheck__) |
| 301 | 6.191e-05 | 2.057e-07 | 6.191e-05 | 2.057e-07 | ~:0(<method 'write' of '_io.TextIOWrapper' objects>) |
| 428 | 6.187e-05 | 1.446e-07 | 6.187e-05 | 1.446e-07 | partial_eval.py:1644(<genexpr>) |
| 31 | 6.175e-05 | 1.992e-06 | 0.0002976 | 9.601e-06 | core.py:1534(join) |
| 64 | 6.166e-05 | 9.634e-07 | 0.0003499 | 5.468e-06 | base.py:310(current_name) |
| 114 | 6.126e-05 | 5.374e-07 | 0.0004854 | 4.258e-06 | partial_eval.py:924(<listcomp>) |
| 64 | 6.12e-05 | 9.562e-07 | 0.0002007 | 3.136e-06 | data_structures.py:142(_copy_structure) |
| 198 | 6.116e-05 | 3.089e-07 | 9.761e-05 | 4.93e-07 | utils.py:77(<lambda>) |
| 48 | 6.097e-05 | 1.27e-06 | 0.0005203 | 1.084e-05 | api.py:2895(named_call) |
| 8 | 6.066e-05 | 7.583e-06 | 0.0001712 | 2.14e-05 | base.py:103(create) |
| 254 | 6.055e-05 | 2.384e-07 | 8.771e-05 | 3.453e-07 | partial_eval.py:908(type_substitute) |
| 24 | 6.054e-05 | 2.522e-06 | 0.0003479 | 1.449e-05 | api.py:763(_check_input_dtype_revderiv) |
| 150 | 6.038e-05 | 4.025e-07 | 0.0001475 | 9.836e-07 | utils.py:87(<genexpr>) |
| 4 | 6.014e-05 | 1.503e-05 | 0.0001375 | 3.437e-05 | _stablehlo_ops_gen.py:2607(__init__) |
| 136 | 5.973e-05 | 4.392e-07 | 5.973e-05 | 4.392e-07 | getlimits.py:692(max) |
| 48 | 5.941e-05 | 1.238e-06 | 8.275e-05 | 1.724e-06 | typing.py:1251(__call__) |
| 258 | 5.902e-05 | 2.287e-07 | 5.902e-05 | 2.287e-07 | ad.py:310(<genexpr>) |
| 15 | 5.876e-05 | 3.918e-06 | 0.001105 | 7.367e-05 | lax.py:1174(full) |
| 3 | 5.87e-05 | 1.957e-05 | 0.0001408 | 4.693e-05 | _stablehlo_ops_gen.py:4393(__init__) |
| 368 | 5.869e-05 | 1.595e-07 | 5.869e-05 | 1.595e-07 | ad.py:275(is_undefined_primal) |
| 352 | 5.835e-05 | 1.658e-07 | 5.835e-05 | 1.658e-07 | ~:0(<method 'item' of 'numpy.ndarray' objects>) |
| 12 | 5.832e-05 | 4.86e-06 | 0.0004562 | 3.802e-05 | mlir.py:1328(reshape) |
| 8 | 5.829e-05 | 7.286e-06 | 0.08874 | 0.01109 | transform.py:343(apply_fn) |
| 21 | 5.811e-05 | 2.767e-06 | 6.84e-05 | 3.257e-06 | sharding.py:97(<dictcomp>) |
| 41 | 5.791e-05 | 1.413e-06 | 9.296e-05 | 2.267e-06 | tree_util.py:727(<listcomp>) |
| 26 | 5.752e-05 | 2.212e-06 | 0.0004992 | 1.92e-05 | lax.py:2026(_add_transpose) |
| 48 | 5.736e-05 | 1.195e-06 | 5.736e-05 | 1.195e-06 | ~:0(<method 'match' of 're.Pattern' objects>) |
| 3 | 5.726e-05 | 1.909e-05 | 0.0001415 | 4.717e-05 | _stablehlo_ops_gen.py:1045(__init__) |
| 31 | 5.676e-05 | 1.831e-06 | 5.676e-05 | 1.831e-06 | ~:0(<built-in method numpy.arange>) |
| 74 | 5.644e-05 | 7.628e-07 | 0.0009559 | 1.292e-05 | core.py:1952(greater_equal_shape) |
| 24 | 5.615e-05 | 2.34e-06 | 0.000197 | 8.209e-06 | lax_numpy.py:121(iscomplexobj) |
| 38 | 5.613e-05 | 1.477e-06 | 5.613e-05 | 1.477e-06 | pprint.py:107(__init__) |
| 186 | 5.606e-05 | 3.014e-07 | 5.606e-05 | 3.014e-07 | core.py:906(__lt__) |
| 2 | 5.537e-05 | 2.769e-05 | 0.3209 | 0.1604 | api.py:720(value_and_grad_f) |
| 18 | 5.496e-05 | 3.054e-06 | 0.000677 | 3.761e-05 | lax.py:741(broadcast) |
| 382 | 5.486e-05 | 1.436e-07 | 5.486e-05 | 1.436e-07 | pjit.py:1775(<genexpr>) |
| 60 | 5.443e-05 | 9.071e-07 | 0.005533 | 9.221e-05 | lax.py:360(integer_pow) |
| 62 | 5.437e-05 | 8.77e-07 | 5.437e-05 | 8.77e-07 | ~:0(<built-in method jaxlib.xla_extension.device_count>) |
| 33 | 5.433e-05 | 1.646e-06 | 5.433e-05 | 1.646e-06 | reductions.py:64(_upcast_f16) |
| 13 | 5.408e-05 | 4.16e-06 | 0.01062 | 0.0008171 | ad.py:535(standard_jvp2) |
| 98 | 5.387e-05 | 5.497e-07 | 5.387e-05 | 5.497e-07 | functools.py:65(wraps) |
| 25 | 5.375e-05 | 2.15e-06 | 0.00151 | 6.039e-05 | mlir.py:304(_python_scalar_handler) |
| 62 | 5.364e-05 | 8.651e-07 | 9.208e-05 | 1.485e-06 | function_base.py:2410(<listcomp>) |
| 54 | 5.344e-05 | 9.896e-07 | 0.0001293 | 2.395e-06 | utils.py:109(indent) |
| 103 | 5.297e-05 | 5.142e-07 | 5.297e-05 | 5.142e-07 | lax.py:3556(<genexpr>) |
| 493 | 5.274e-05 | 1.07e-07 | 5.274e-05 | 1.07e-07 | core.py:627(_assert_live) |
| 128 | 5.257e-05 | 4.107e-07 | 0.0002758 | 2.155e-06 | base.py:305(current_module) |
| 128 | 5.244e-05 | 4.097e-07 | 0.000854 | 6.672e-06 | partial_eval.py:149(lift) |
| 128 | 5.238e-05 | 4.093e-07 | 0.0001187 | 9.275e-07 | base.py:347(assert_context) |
| 205 | 5.232e-05 | 2.552e-07 | 8.482e-05 | 4.138e-07 | pxla.py:2434(<genexpr>) |
| 37 | 5.176e-05 | 1.399e-06 | 0.00424 | 0.0001146 | lax.py:1078(_reduce_sum) |
| 38 | 5.134e-05 | 1.351e-06 | 0.000513 | 1.35e-05 | pprint.py:156(pformat) |
| 1 | 5.116e-05 | 5.116e-05 | 0.08713 | 0.08713 | algorithm.py:212(compute_loss_q) |
| 100 | 5.064e-05 | 5.064e-07 | 5.618e-05 | 5.618e-07 | lax.py:4610(ranges_like) |
| 4 | 5.036e-05 | 1.259e-05 | 0.0005977 | 0.0001494 | ad.py:73(<listcomp>) |
| 336 | 4.992e-05 | 1.486e-07 | 4.992e-05 | 1.486e-07 | lax.py:2424(<genexpr>) |
| 48/16 | 4.97e-05 | 3.106e-06 | 0.0006592 | 4.12e-05 | data_structures.py:150(<dictcomp>) |
| 110 | 4.965e-05 | 4.514e-07 | 0.0001885 | 1.713e-06 | pjit.py:741(<lambda>) |
| 253 | 4.961e-05 | 1.961e-07 | 4.961e-05 | 1.961e-07 | partial_eval.py:1395(<genexpr>) |
| 94/38 | 4.945e-05 | 1.301e-06 | 0.0003267 | 8.598e-06 | pprint.py:463(format) |
| 118 | 4.927e-05 | 4.175e-07 | 5.871e-05 | 4.976e-07 | util.py:248(<genexpr>) |
| 122 | 4.927e-05 | 4.038e-07 | 6.588e-05 | 5.4e-07 | partial_eval.py:707(full_lower) |
| 96 | 4.89e-05 | 5.094e-07 | 4.89e-05 | 5.094e-07 | contextlib.py:460(__init__) |
| 64 | 4.861e-05 | 7.595e-07 | 4.861e-05 | 7.595e-07 | pxla.py:2422(<listcomp>) |
| 16 | 4.839e-05 | 3.025e-06 | 0.0003519 | 2.199e-05 | xla.py:178(_canonicalize_python_scalar_dtype) |
| 2 | 4.836e-05 | 2.418e-05 | 0.0004389 | 0.0002195 | hlo_helpers.py:23(custom_call) |
| 2 | 4.81e-05 | 2.405e-05 | 0.0004945 | 0.0002472 | gpu_prng.py:45(_threefry2x32_lowering) |
| 15 | 4.785e-05 | 3.19e-06 | 0.001009 | 6.724e-05 | source_info_util.py:171(summarize) |
| 346 | 4.778e-05 | 1.381e-07 | 4.778e-05 | 1.381e-07 | ~:0(<built-in method builtins.setattr>) |
| 321 | 4.752e-05 | 1.48e-07 | 4.752e-05 | 1.48e-07 | lax.py:2921(<genexpr>) |
| 12 | 4.741e-05 | 3.95e-06 | 0.003019 | 0.0002515 | lax.py:2754(_broadcast_in_dim_jvp_rule) |
| 8 | 4.717e-05 | 5.896e-06 | 0.0008458 | 0.0001057 | base.py:215(new_context) |
| 15 | 4.706e-05 | 3.137e-06 | 0.0005913 | 3.942e-05 | _func_ops_ext.py:236(__init__) |
| 62 | 4.691e-05 | 7.566e-07 | 9.529e-05 | 1.537e-06 | __init__.py:1467(debug) |
| 15 | 4.686e-05 | 3.124e-06 | 8.554e-05 | 5.703e-06 | pjit.py:1445(<listcomp>) |
| 55 | 4.68e-05 | 8.509e-07 | 0.001016 | 1.848e-05 | <__array_function__ internals>:177(delete) |
| 8 | 4.668e-05 | 5.835e-06 | 0.0001515 | 1.894e-05 | core.py:1544(str_short) |
| 318 | 4.646e-05 | 1.461e-07 | 4.646e-05 | 1.461e-07 | ad.py:270(__init__) |
| 44 | 4.631e-05 | 1.052e-06 | 8.99e-05 | 2.043e-06 | pjit.py:1228(<listcomp>) |
| 13 | 4.61e-05 | 3.546e-06 | 0.009461 | 0.0007278 | lax.py:1747(<lambda>) |
| 224 | 4.585e-05 | 2.047e-07 | 4.585e-05 | 2.047e-07 | utils.py:71(<genexpr>) |
| 2 | 4.566e-05 | 2.283e-05 | 0.1677 | 0.08386 | pjit.py:1533(_pjit_jvp) |
| 39 | 4.559e-05 | 1.169e-06 | 0.0004785 | 1.227e-05 | lax.py:1552(<listcomp>) |
| 32 | 4.558e-05 | 1.424e-06 | 0.0008131 | 2.541e-05 | basic.py:132(__init__) |
| 62 | 4.532e-05 | 7.31e-07 | 5.863e-05 | 9.456e-07 | function_base.py:346(iterable) |
| 44 | 4.529e-05 | 1.029e-06 | 4.529e-05 | 1.029e-06 | pjit.py:1162(<listcomp>) |
| 63 | 4.511e-05 | 7.16e-07 | 0.0084 | 0.0001333 | lax.py:400(add) |
| 2 | 4.482e-05 | 2.241e-05 | 0.005102 | 0.002551 | prng.py:1060(threefry_2x32) |
| 186 | 4.46e-05 | 2.398e-07 | 4.46e-05 | 2.398e-07 | effects.py:63(<listcomp>) |
| 62 | 4.42e-05 | 7.13e-07 | 4.42e-05 | 7.13e-07 | ~:0(<built-in method jaxlib.xla_extension.replica_count>) |
| 64 | 4.389e-05 | 6.858e-07 | 6.267e-05 | 9.792e-07 | base.py:978(assert_jax_usage) |
| 77 | 4.381e-05 | 5.69e-07 | 4.381e-05 | 5.69e-07 | mlir.py:959(<listcomp>) |
| 92 | 4.339e-05 | 4.717e-07 | 5.188e-05 | 5.639e-07 | lax.py:3076(<genexpr>) |
| 8 | 4.338e-05 | 5.423e-06 | 0.06067 | 0.007584 | reductions.py:487(_std) |
| 4 | 4.338e-05 | 1.084e-05 | 0.0003757 | 9.393e-05 | partial_eval.py:793(<listcomp>) |
| 11 | 4.322e-05 | 3.929e-06 | 0.0001324 | 1.203e-05 | xla.py:211(_make_shaped_array_for_numpy_scalar) |
| 12 | 4.274e-05 | 3.562e-06 | 6.402e-05 | 5.335e-06 | lax.py:3218(_transpose_shape_rule) |
| 138 | 4.262e-05 | 3.088e-07 | 6.059e-05 | 4.391e-07 | ad.py:220(write_primal) |
| 8 | 4.258e-05 | 5.323e-06 | 0.001731 | 0.0002163 | slicing.py:643(slice_in_dim) |
| 96 | 4.216e-05 | 4.391e-07 | 6.869e-05 | 7.155e-07 | <string>:1(<lambda>) |
| 6 | 4.206e-05 | 7.01e-06 | 0.0003297 | 5.494e-05 | prng.py:837(_threefry2x32_abstract_eval) |
| 12 | 4.196e-05 | 3.496e-06 | 0.0002222 | 1.852e-05 | _stablehlo_ops_gen.py:3928(__init__) |
| 8 | 4.186e-05 | 5.232e-06 | 0.08878 | 0.0111 | transform.py:120(apply_fn) |
| 25 | 4.154e-05 | 1.661e-06 | 4.154e-05 | 1.661e-06 | ~:0(<method 'take' of 'numpy.ndarray' objects>) |
| 2 | 4.148e-05 | 2.074e-05 | 0.02988 | 0.01494 | pjit.py:1722(_pjit_transpose) |
| 52 | 4.125e-05 | 7.933e-07 | 0.01208 | 0.0002322 | core.py:661(__pow__) |
| 116 | 4.12e-05 | 3.552e-07 | 6.006e-05 | 5.177e-07 | api_util.py:398(is_hashable) |
| 84 | 4.108e-05 | 4.89e-07 | 4.108e-05 | 4.89e-07 | _dtype.py:24(_kind_name) |
| 50 | 4.095e-05 | 8.189e-07 | 0.001017 | 2.034e-05 | <__array_function__ internals>:177(argsort) |
| 129 | 4.08e-05 | 3.163e-07 | 6.513e-05 | 5.049e-07 | pjit.py:909(<genexpr>) |
| 8 | 4.068e-05 | 5.085e-06 | 0.0002519 | 3.149e-05 | core.py:720(_pretty_print) |
| 32 | 4.064e-05 | 1.27e-06 | 0.006785 | 0.000212 | lax_numpy.py:1113(broadcast_to) |
| 17 | 4.055e-05 | 2.385e-06 | 0.003928 | 0.000231 | lax.py:2013(_add_jvp) |
| 210 | 4.047e-05 | 1.927e-07 | 4.047e-05 | 1.927e-07 | api_util.py:295(<genexpr>) |
| 54 | 4.009e-05 | 7.424e-07 | 0.0008549 | 1.583e-05 | utils.py:94(<genexpr>) |
| 13 | 3.985e-05 | 3.065e-06 | 0.0005296 | 4.074e-05 | mlir.py:1419(add_jaxvals_lowering) |
| 64 | 3.98e-05 | 6.218e-07 | 3.98e-05 | 6.218e-07 | pxla.py:2421(<listcomp>) |
| 3 | 3.971e-05 | 1.324e-05 | 0.02782 | 0.009273 | core.py:35(__call__) |
| 392 | 3.948e-05 | 1.007e-07 | 3.948e-05 | 1.007e-07 | partial_eval.py:930(<genexpr>) |
| 47 | 3.944e-05 | 8.392e-07 | 0.006522 | 0.0001388 | lax.py:412(div) |
| 112 | 3.932e-05 | 3.51e-07 | 4.93e-05 | 4.402e-07 | source_info_util.py:90(__len__) |
| 62 | 3.921e-05 | 6.325e-07 | 4.494e-05 | 7.248e-07 | xla_bridge.py:357(is_known_platform) |
| 238 | 3.919e-05 | 1.647e-07 | 3.919e-05 | 1.647e-07 | partial_eval.py:220(<genexpr>) |
| 53 | 3.916e-05 | 7.388e-07 | 4.416e-05 | 8.332e-07 | numeric.py:1878(isscalar) |
| 321 | 3.911e-05 | 1.218e-07 | 3.911e-05 | 1.218e-07 | lax.py:2204(_convert_element_type_weak_type_rule) |
| 50 | 3.906e-05 | 7.811e-07 | 0.0009314 | 1.863e-05 | fromnumeric.py:1038(argsort) |
| 38 | 3.904e-05 | 1.027e-06 | 0.0003717 | 9.783e-06 | pprint.py:454(_repr) |
| 64 | 3.903e-05 | 6.098e-07 | 3.903e-05 | 6.098e-07 | pxla.py:1690(__init__) |
| 2 | 3.897e-05 | 1.949e-05 | 0.05454 | 0.02727 | partial_eval.py:1066(fun) |
| 4 | 3.864e-05 | 9.661e-06 | 6.107e-05 | 1.527e-05 | partial_eval.py:792(<listcomp>) |
| 2 | 3.814e-05 | 1.907e-05 | 0.000204 | 0.000102 | partial_eval.py:1069(<listcomp>) |
| 321 | 3.794e-05 | 1.182e-07 | 3.794e-05 | 1.182e-07 | lax.py:2926(<genexpr>) |
| 72 | 3.793e-05 | 5.268e-07 | 3.793e-05 | 5.268e-07 | api_util.py:349(<dictcomp>) |
| 5 | 3.772e-05 | 7.543e-06 | 0.003003 | 0.0006006 | lax_numpy.py:850(squeeze) |
| 70 | 3.764e-05 | 5.377e-07 | 0.00213 | 3.042e-05 | util.py:355(_check_no_float0s) |
| 404 | 3.737e-05 | 9.249e-08 | 3.737e-05 | 9.249e-08 | partial_eval.py:1679(<genexpr>) |
| 112 | 3.722e-05 | 3.323e-07 | 3.722e-05 | 3.323e-07 | config.py:165(get_config) |
| 62 | 3.698e-05 | 5.965e-07 | 3.698e-05 | 5.965e-07 | mlir.py:927(<listcomp>) |
| 210 | 3.687e-05 | 1.756e-07 | 3.687e-05 | 1.756e-07 | pjit.py:1586(<genexpr>) |
| 23 | 3.65e-05 | 1.587e-06 | 0.005648 | 0.0002456 | dispatch.py:202(prim_fun) |
| 136 | 3.646e-05 | 2.681e-07 | 3.646e-05 | 2.681e-07 | utils.py:30(compare_or_false) |
| 4 | 3.646e-05 | 9.116e-06 | 5.603e-05 | 1.401e-05 | util.py:214(_remove_duplicates) |
| 24 | 3.621e-05 | 1.509e-06 | 0.004999 | 0.0002083 | transform.py:96(orderth_norm) |
| 13 | 3.604e-05 | 2.772e-06 | 0.002944 | 0.0002264 | pjit.py:1761(_dce_jaxpr_pjit) |
| 22 | 3.602e-05 | 1.637e-06 | 3.602e-05 | 1.637e-06 | ~:0(<method 'item' of 'numpy.generic' objects>) |
| 31 | 3.589e-05 | 1.158e-06 | 0.00318 | 0.0001026 | util.py:252(<listcomp>) |
| 8 | 3.585e-05 | 4.482e-06 | 0.002701 | 0.0003376 | slicing.py:672(index_in_dim) |
| 8 | 3.584e-05 | 4.48e-06 | 0.0007882 | 9.852e-05 | base.py:851(__init__) |
| 2 | 3.571e-05 | 1.785e-05 | 8.197e-05 | 4.098e-05 | _stablehlo_ops_gen.py:3074(__init__) |
| 8 | 3.569e-05 | 4.461e-06 | 0.0002641 | 3.302e-05 | pretty_printer.py:61(format) |
| 96 | 3.537e-05 | 3.684e-07 | 3.537e-05 | 3.684e-07 | ~:0(<built-in method sys.exc_info>) |
| 48 | 3.521e-05 | 7.335e-07 | 9.257e-05 | 1.928e-06 | module.py:446(<lambda>) |
| 56 | 3.507e-05 | 6.263e-07 | 0.0001071 | 1.912e-06 | util.py:414(<genexpr>) |
| 2 | 3.506e-05 | 1.753e-05 | 8.047e-05 | 4.024e-05 | _stablehlo_ops_gen.py:925(__init__) |
| 19 | 3.505e-05 | 1.845e-06 | 8.307e-05 | 4.372e-06 | tree_util.py:592(<lambda>) |
| 92 | 3.473e-05 | 3.775e-07 | 3.862e-05 | 4.198e-07 | lax.py:3074(<genexpr>) |
| 35 | 3.473e-05 | 9.923e-07 | 4.111e-05 | 1.175e-06 | dtypes.py:87(to_inexact_dtype) |
| 2 | 3.454e-05 | 1.727e-05 | 0.2893 | 0.1447 | ad.py:116(linearize) |
| 1 | 3.452e-05 | 3.452e-05 | 0.006172 | 0.006172 | random.py:270(_uniform) |
| 12 | 3.426e-05 | 2.855e-06 | 0.0007551 | 6.292e-05 | ad.py:71(jvpfun) |
| 30 | 3.403e-05 | 1.134e-06 | 3.949e-05 | 1.316e-06 | mlir.py:827(effects) |
| 40 | 3.403e-05 | 8.506e-07 | 0.03083 | 0.0007709 | core.py:653(__truediv__) |
| 27 | 3.396e-05 | 1.258e-06 | 0.0001563 | 5.788e-06 | fromnumeric.py:3194(size) |
| 25 | 3.344e-05 | 1.338e-06 | 3.706e-05 | 1.482e-06 | core.py:1999(dimension_as_value) |
| 15 | 3.337e-05 | 2.224e-06 | 0.0008824 | 5.883e-05 | source_info_util.py:173(<listcomp>) |
| 48 | 3.323e-05 | 6.923e-07 | 3.323e-05 | 6.923e-07 | utils.py:84(<genexpr>) |
| 77 | 3.321e-05 | 4.312e-07 | 3.321e-05 | 4.312e-07 | partial_eval.py:1389(<listcomp>) |
| 1 | 3.303e-05 | 3.303e-05 | 0.0001021 | 0.0001021 | seeding.py:9(np_random) |
| 321 | 3.3e-05 | 1.028e-07 | 3.3e-05 | 1.028e-07 | lax.py:2201(_convert_element_type_dtype_rule) |
| 26 | 3.298e-05 | 1.269e-06 | 0.001274 | 4.899e-05 | util.py:294(<listcomp>) |
| 48 | 3.281e-05 | 6.835e-07 | 6.774e-05 | 1.411e-06 | contextlib.py:118(_recreate_cm) |
| 29 | 3.262e-05 | 1.125e-06 | 5.249e-05 | 1.81e-06 | core.py:1510(<lambda>) |
| 80 | 3.248e-05 | 4.06e-07 | 3.248e-05 | 4.06e-07 | lax.py:3080(<genexpr>) |
| 14 | 3.24e-05 | 2.314e-06 | 6.927e-05 | 4.948e-06 | pjit.py:1585(keep_where) |
| 8 | 3.228e-05 | 4.035e-06 | 4.303e-05 | 5.379e-06 | base.py:207(__exit__) |
| 37 | 3.203e-05 | 8.658e-07 | 3.203e-05 | 8.658e-07 | slicing.py:762(<listcomp>) |
| 28 | 3.194e-05 | 1.141e-06 | 0.002735 | 9.767e-05 | util.py:249(<listcomp>) |
| 112 | 3.188e-05 | 2.846e-07 | 3.188e-05 | 2.846e-07 | ad.py:282(get_primitive_transpose) |
| 63 | 3.176e-05 | 5.042e-07 | 4.576e-05 | 7.263e-07 | tree_util.py:592(<genexpr>) |
| 77 | 3.136e-05 | 4.073e-07 | 3.136e-05 | 4.073e-07 | mlir.py:997(<listcomp>) |
| 87 | 3.133e-05 | 3.601e-07 | 0.0003594 | 4.131e-06 | slicing.py:752(<genexpr>) |
| 8 | 3.133e-05 | 3.916e-06 | 0.04001 | 0.005001 | ad.py:496(linear_jvp) |
| 8 | 3.106e-05 | 3.883e-06 | 0.0002184 | 2.73e-05 | base.py:195(__enter__) |
| 23 | 3.106e-05 | 1.351e-06 | 3.106e-05 | 1.351e-06 | dispatch.py:238(<listcomp>) |
| 360 | 3.066e-05 | 8.516e-08 | 3.066e-05 | 8.516e-08 | ~:0(<method 'append' of 'collections.deque' objects>) |
| 3 | 3.054e-05 | 1.018e-05 | 0.008759 | 0.00292 | numerics.py:103(safe_int32_increment) |
| 27 | 3.049e-05 | 1.129e-06 | 0.0002506 | 9.282e-06 | <__array_function__ internals>:177(size) |
| 25 | 3.041e-05 | 1.216e-06 | 0.002127 | 8.509e-05 | reductions.py:183(_cast_to_numeric) |
| 62 | 3.015e-05 | 4.863e-07 | 3.015e-05 | 4.863e-07 | compilation_cache.py:263(is_initialized) |
| 2 | 3.006e-05 | 1.503e-05 | 0.09642 | 0.04821 | ad.py:689(jvp_jaxpr) |
| 19 | 3e-05 | 1.579e-06 | 3.288e-05 | 1.731e-06 | sharding_impls.py:354(devices_indices_map) |
| 159 | 2.996e-05 | 1.884e-07 | 2.996e-05 | 1.884e-07 | pjit.py:880(<genexpr>) |
| 2 | 2.965e-05 | 1.482e-05 | 0.05996 | 0.02998 | partial_eval.py:998(partial_eval_jaxpr_nounits) |
| 31 | 2.936e-05 | 9.472e-07 | 0.0003302 | 1.065e-05 | core.py:1311(lattice_join) |
| 26 | 2.897e-05 | 1.114e-06 | 0.000296 | 1.138e-05 | ad.py:124(<genexpr>) |
| 5 | 2.873e-05 | 5.746e-06 | 0.001883 | 0.0003767 | lax.py:657(dot) |
| 168 | 2.824e-05 | 1.681e-07 | 2.824e-05 | 1.681e-07 | lax.py:2417(<genexpr>) |
| 148 | 2.772e-05 | 1.873e-07 | 2.772e-05 | 1.873e-07 | util.py:284(__init__) |
| 165 | 2.766e-05 | 1.676e-07 | 2.766e-05 | 1.676e-07 | core.py:2551(<listcomp>) |
| 175 | 2.728e-05 | 1.559e-07 | 2.728e-05 | 1.559e-07 | ad.py:195(<genexpr>) |
| 2 | 2.724e-05 | 1.362e-05 | 4.419e-05 | 2.209e-05 | pjit.py:1622(<listcomp>) |
| 7 | 2.707e-05 | 3.867e-06 | 0.0004251 | 6.073e-05 | lax.py:792(reshape) |
| 49 | 2.702e-05 | 5.515e-07 | 7.43e-05 | 1.516e-06 | pjit.py:1737(<genexpr>) |
| 79 | 2.695e-05 | 3.412e-07 | 2.695e-05 | 3.412e-07 | mlir.py:510(set_tokens_out) |
| 8 | 2.688e-05 | 3.36e-06 | 0.000233 | 2.912e-05 | reductions.py:459(_var_promote_types) |
| 2 | 2.674e-05 | 1.337e-05 | 0.2895 | 0.1447 | api.py:2252(_vjp) |
| 16 | 2.67e-05 | 1.669e-06 | 3.581e-05 | 2.238e-06 | module.py:508(next_module_number) |
| 5 | 2.664e-05 | 5.328e-06 | 0.05359 | 0.01072 | core.py:105(<lambda>) |
| 139 | 2.658e-05 | 1.912e-07 | 2.658e-05 | 1.912e-07 | mlir.py:1134(<genexpr>) |
| 177 | 2.655e-05 | 1.5e-07 | 2.655e-05 | 1.5e-07 | partial_eval.py:1396(<genexpr>) |
| 2 | 2.649e-05 | 1.325e-05 | 0.001023 | 0.0005115 | lax_numpy.py:1120(_split) |
| 48 | 2.649e-05 | 5.518e-07 | 0.0001191 | 2.48e-06 | module.py:557(check_name) |
| 12 | 2.649e-05 | 2.207e-06 | 0.005435 | 0.0004529 | lax.py:2472(_dot_general_transpose_rhs) |
| 12 | 2.643e-05 | 2.203e-06 | 0.005465 | 0.0004554 | partial_eval.py:777(trace_to_subjaxpr_nounits) |
| 3 | 2.619e-05 | 8.73e-06 | 0.0001674 | 5.579e-05 | lax.py:3338(_select_hlo_lowering) |
| 3 | 2.608e-05 | 8.693e-06 | 0.0002627 | 8.757e-05 | mlir.py:1426(compare_hlo) |
| 312 | 2.599e-05 | 8.329e-08 | 2.599e-05 | 8.329e-08 | ~:0(<method 'pop' of 'collections.deque' objects>) |
| 78 | 2.578e-05 | 3.305e-07 | 2.578e-05 | 3.305e-07 | ad.py:465(full_lower) |
| 5 | 2.574e-05 | 5.148e-06 | 0.04565 | 0.00913 | core.py:47(__call__) |
| 92 | 2.568e-05 | 2.791e-07 | 3.505e-05 | 3.809e-07 | ad_util.py:83(<lambda>) |
| 39 | 2.562e-05 | 6.568e-07 | 9.348e-05 | 2.397e-06 | lax.py:1052(_get_sum_identity) |
| 130 | 2.558e-05 | 1.968e-07 | 2.558e-05 | 1.968e-07 | utils.py:112(<genexpr>) |
| 16 | 2.553e-05 | 1.596e-06 | 0.006609 | 0.000413 | ufuncs.py:251(true_divide) |
| 4/2 | 2.55e-05 | 1.275e-05 | 0.2889 | 0.1444 | partial_eval.py:763(trace_to_jaxpr_nounits) |
| 4 | 2.546e-05 | 6.365e-06 | 3.846e-05 | 9.615e-06 | partial_eval.py:791(<listcomp>) |
| 112 | 2.515e-05 | 2.245e-07 | 2.515e-05 | 2.245e-07 | __init__.py:658(update) |
| 56 | 2.504e-05 | 4.472e-07 | 3.001e-05 | 5.359e-07 | util.py:419(<genexpr>) |
| 8 | 2.496e-05 | 3.12e-06 | 8.62e-05 | 1.078e-05 | base.py:165(__init__) |
| 62 | 2.487e-05 | 4.011e-07 | 2.487e-05 | 4.011e-07 | ~:0(<built-in method jaxlib.xla_extension.computation_count>) |
| 50 | 2.461e-05 | 4.922e-07 | 2.461e-05 | 4.922e-07 | lax.py:4620(<listcomp>) |
| 57 | 2.43e-05 | 4.263e-07 | 0.001563 | 2.742e-05 | ad.py:542(add_tangents) |
| 8 | 2.427e-05 | 3.034e-06 | 3.218e-05 | 4.022e-06 | base.py:71(current) |
| 5 | 2.419e-05 | 4.838e-06 | 0.003083 | 0.0006166 | lax_numpy.py:2920(dot) |
| 28 | 2.414e-05 | 8.62e-07 | 0.0001161 | 4.148e-06 | <__array_function__ internals>:177(where) |
| 77 | 2.375e-05 | 3.085e-07 | 2.375e-05 | 3.085e-07 | pjit.py:1726(<genexpr>) |
| 2 | 2.365e-05 | 1.183e-05 | 0.0008815 | 0.0004408 | prng.py:924(_threefry2x32_gpu_lowering) |
| 24 | 2.363e-05 | 9.848e-07 | 0.01383 | 0.0005763 | transform.py:518(<lambda>) |
| 62 | 2.352e-05 | 3.793e-07 | 2.352e-05 | 3.793e-07 | mlir.py:1207(lower_fun) |
| 6 | 2.338e-05 | 3.896e-06 | 6.155e-05 | 1.026e-05 | util.py:128(merge_lists) |
| 48 | 2.334e-05 | 4.862e-07 | 2.334e-05 | 4.862e-07 | data_structures.py:43(__init__) |
| 26 | 2.32e-05 | 8.922e-07 | 0.009484 | 0.0003648 | ad.py:537(<genexpr>) |
| 15 | 2.319e-05 | 1.546e-06 | 0.000258 | 1.72e-05 | mlir.py:794(token_type) |
| 6 | 2.314e-05 | 3.856e-06 | 3.367e-05 | 5.612e-06 | dtypes.py:511(_least_upper_bound) |
| 188 | 2.314e-05 | 1.231e-07 | 2.314e-05 | 1.231e-07 | lax.py:170(_identity) |
| 169 | 2.307e-05 | 1.365e-07 | 2.307e-05 | 1.365e-07 | tree_util.py:242(<lambda>) |
| 28 | 2.298e-05 | 8.207e-07 | 0.003838 | 0.0001371 | lax.py:404(sub) |
| 15 | 2.277e-05 | 1.518e-06 | 0.0001658 | 1.106e-05 | mlir.py:550(flatten_lowering_ir_args) |
| 2 | 2.275e-05 | 1.138e-05 | 6.981e-05 | 3.49e-05 | box.py:224(contains) |
| 2 | 2.266e-05 | 1.133e-05 | 0.09636 | 0.04818 | ad.py:696(_jvp_jaxpr) |
| 75 | 2.251e-05 | 3.001e-07 | 3.012e-05 | 4.015e-07 | lax.py:934(<genexpr>) |
| 8 | 2.25e-05 | 2.812e-06 | 2.667e-05 | 3.334e-06 | _parser.py:1090(expand_template) |
| 15 | 2.241e-05 | 1.494e-06 | 4.676e-05 | 3.117e-06 | mlir.py:834(update_tokens) |
| 64 | 2.238e-05 | 3.497e-07 | 2.238e-05 | 3.497e-07 | dispatch.py:373(<setcomp>) |
| 94 | 2.227e-05 | 2.369e-07 | 3.088e-05 | 3.285e-07 | pjit.py:1577(<genexpr>) |
| 168 | 2.201e-05 | 1.31e-07 | 2.201e-05 | 1.31e-07 | lax.py:2434(<genexpr>) |
| 6 | 2.194e-05 | 3.657e-06 | 3.039e-05 | 5.064e-06 | util.py:131(<listcomp>) |
| 101 | 2.192e-05 | 2.17e-07 | 2.192e-05 | 2.17e-07 | util.py:296(__init__) |
| 2 | 2.189e-05 | 1.095e-05 | 0.0001009 | 5.044e-05 | partial_eval.py:1448(move_binders_to_front) |
| 2 | 2.165e-05 | 1.083e-05 | 0.03026 | 0.01513 | ad.py:143(unbound_vjp) |
| 3 | 2.15e-05 | 7.165e-06 | 0.003029 | 0.00101 | lax_numpy.py:1006(where) |
| 116 | 2.143e-05 | 1.847e-07 | 2.143e-05 | 1.847e-07 | api_util.py:156(__init__) |
| 2 | 2.137e-05 | 1.069e-05 | 8.757e-05 | 4.378e-05 | hlo_helpers.py:55(<listcomp>) |
| 2 | 2.124e-05 | 1.062e-05 | 6.566e-05 | 3.283e-05 | partial_eval.py:1076(<listcomp>) |
| 31 | 2.114e-05 | 6.82e-07 | 2.345e-05 | 7.563e-07 | util.py:250(<setcomp>) |
| 5 | 2.111e-05 | 4.221e-06 | 0.006519 | 0.001304 | core.py:41(__init__) |
| 10 | 2.111e-05 | 2.111e-06 | 0.0008334 | 8.334e-05 | slicing.py:819(_slice_lower) |
| 5 | 2.103e-05 | 4.206e-06 | 6.348e-05 | 1.27e-05 | api_util.py:576(_np_scalar_abstractify) |
| 8 | 2.098e-05 | 2.622e-06 | 5.764e-05 | 7.205e-06 | pretty_printer.py:255(_align_annotations) |
| 100 | 2.094e-05 | 2.094e-07 | 2.094e-05 | 2.094e-07 | ~:0(<method 'update' of 'dict' objects>) |
| 30 | 2.083e-05 | 6.942e-07 | 7.21e-05 | 2.403e-06 | lax.py:4623(canonicalize_precision) |
| 3 | 2.077e-05 | 6.924e-06 | 0.000414 | 0.000138 | lax.py:3517(_reduce_sum_transpose_rule) |
| 168 | 2.076e-05 | 1.236e-07 | 2.076e-05 | 1.236e-07 | lax.py:2418(<genexpr>) |
| 62 | 2.071e-05 | 3.34e-07 | 2.071e-05 | 3.34e-07 | ~:0(<method 'getvalue' of '_io.BytesIO' objects>) |
| 3 | 2.064e-05 | 6.88e-06 | 5.474e-05 | 1.825e-05 | ad.py:95(<listcomp>) |
| 2 | 2.06e-05 | 1.03e-05 | 0.0001586 | 7.929e-05 | pxla.py:3143(from_trivial_jaxpr) |
| 4 | 2.057e-05 | 5.143e-06 | 0.001711 | 0.0004278 | reductions.py:248(_reduce_max) |
| 15 | 2.047e-05 | 1.365e-06 | 3.631e-05 | 2.421e-06 | pjit.py:1447(<listcomp>) |
| 12 | 2.043e-05 | 1.702e-06 | 7.033e-05 | 5.861e-06 | numeric.py:150(ones) |
| 30 | 2.032e-05 | 6.772e-07 | 2.032e-05 | 6.772e-07 | ~:0(<built-in method math.prod>) |
| 25 | 2.029e-05 | 8.116e-07 | 0.0001536 | 6.145e-06 | fromnumeric.py:93(take) |
| 4 | 2.023e-05 | 5.057e-06 | 0.001677 | 0.0004192 | reductions.py:263(_reduce_min) |
| 2 | 2.017e-05 | 1.008e-05 | 0.0001014 | 5.07e-05 | ad.py:712(<listcomp>) |
| 16 | 2.016e-05 | 1.26e-06 | 0.0009988 | 6.243e-05 | ad.py:514(linear_transpose2) |
| 25 | 2.005e-05 | 8.021e-07 | 0.0001932 | 7.727e-06 | <__array_function__ internals>:177(take) |
| 1 | 2.004e-05 | 2.004e-05 | 0.0003996 | 0.0003996 | passive_env_checker.py:211(env_step_passive_checker) |
| 2 | 1.981e-05 | 9.906e-06 | 0.2893 | 0.1447 | ad.py:137(vjp) |
| 84 | 1.968e-05 | 2.343e-07 | 1.968e-05 | 2.343e-07 | pjit.py:1541(<genexpr>) |
| 15 | 1.963e-05 | 1.308e-06 | 1.963e-05 | 1.308e-06 | source_info_util.py:164(_summarize_frame) |
| 1 | 1.962e-05 | 1.962e-05 | 0.0002756 | 0.0002756 | lax.py:2929(_concatenate_transpose_rule) |
| 31 | 1.961e-05 | 6.327e-07 | 0.0003498 | 1.128e-05 | ad_util.py:46(add_abstract) |
| 4 | 1.951e-05 | 4.878e-06 | 0.1394 | 0.03486 | core.py:226(jaxpr_as_fun) |
| 155 | 1.946e-05 | 1.255e-07 | 1.946e-05 | 1.255e-07 | dispatch.py:338(<genexpr>) |
| 10 | 1.939e-05 | 1.939e-06 | 0.002434 | 0.0002434 | ufuncs.py:60(<lambda>) |
| 62 | 1.927e-05 | 3.108e-07 | 1.927e-05 | 3.108e-07 | ~:0(<method '__exit__' of '_thread.lock' objects>) |
| 16 | 1.92e-05 | 1.2e-06 | 0.0001589 | 9.934e-06 | utils.py:256(<lambda>) |
| 6 | 1.918e-05 | 3.197e-06 | 6.072e-05 | 1.012e-05 | pjit.py:1725(prune_type) |
| 15 | 1.917e-05 | 1.278e-06 | 6.74e-05 | 4.493e-06 | pjit.py:1432(_pjit_abstract_eval) |
| 64 | 1.913e-05 | 2.988e-07 | 1.913e-05 | 2.988e-07 | dispatch.py:382(apply_outfeed_rewriter) |
| 4 | 1.911e-05 | 4.778e-06 | 2.594e-05 | 6.485e-06 | partial_eval.py:782(<listcomp>) |
| 50 | 1.911e-05 | 3.821e-07 | 2.727e-05 | 5.453e-07 | inspect.py:300(ismethod) |
| 62 | 1.905e-05 | 3.073e-07 | 1.905e-05 | 3.073e-07 | dispatch.py:277(should_tuple_args) |
| 16 | 1.902e-05 | 1.189e-06 | 1.902e-05 | 1.189e-06 | __init__.py:315(_subx) |
| 8 | 1.898e-05 | 2.372e-06 | 0.003605 | 0.0004506 | ufuncs.py:81(<lambda>) |
| 3 | 1.898e-05 | 6.326e-06 | 0.0003344 | 0.0001115 | lax.py:2160(_compare_lower_hlo) |
| 2 | 1.896e-05 | 9.481e-06 | 0.0001641 | 8.207e-05 | hlo_helpers.py:75(<listcomp>) |
| 29 | 1.884e-05 | 6.497e-07 | 2.425e-05 | 8.363e-07 | util.py:302(__eq__) |
| 2 | 1.883e-05 | 9.413e-06 | 0.0001269 | 6.344e-05 | api.py:670(value_and_grad) |
| 3 | 1.881e-05 | 6.268e-06 | 4.971e-05 | 1.657e-05 | lax.py:1457(_iter) |
| 12 | 1.868e-05 | 1.557e-06 | 0.001254 | 0.0001045 | lax.py:2766(_broadcast_in_dim_partial_eval) |
| 2 | 1.862e-05 | 9.308e-06 | 3.084e-05 | 1.542e-05 | pxla.py:3231(_out_shardings_for_trivial) |
| 4 | 1.855e-05 | 4.638e-06 | 0.0002094 | 5.235e-05 | lax.py:4333(_iota_abstract_eval) |
| 2 | 1.789e-05 | 8.945e-06 | 0.0001082 | 5.409e-05 | lax.py:4373(_iota_lower) |
| 3 | 1.786e-05 | 5.954e-06 | 9.11e-05 | 3.037e-05 | ad.py:102(jvp_subtrace_aux) |
| 30 | 1.786e-05 | 5.952e-07 | 1.786e-05 | 5.952e-07 | prng.py:838(<genexpr>) |
| 3 | 1.785e-05 | 5.95e-06 | 0.0001814 | 6.048e-05 | asserts.py:702(assert_type) |
| 2 | 1.761e-05 | 8.803e-06 | 2.534e-05 | 1.267e-05 | pjit.py:1656(<listcomp>) |
| 1 | 1.76e-05 | 1.76e-05 | 1.76e-05 | 1.76e-05 | ~:0(<function SeedSequence.generate_state at 0x7f4329bef9c0>) |
| 2 | 1.755e-05 | 8.773e-06 | 0.03035 | 0.01518 | api.py:2142(_vjp_pullback_wrapper) |
| 3 | 1.745e-05 | 5.816e-06 | 0.03303 | 0.01101 | core.py:108(<lambda>) |
| 1 | 1.736e-05 | 1.736e-05 | 3.837e-05 | 3.837e-05 | _stablehlo_ops_gen.py:3172(__init__) |
| 128 | 1.736e-05 | 1.356e-07 | 1.736e-05 | 1.356e-07 | base.py:436(check_not_none) |
| 36 | 1.718e-05 | 4.772e-07 | 2.25e-05 | 6.251e-07 | lax_numpy.py:148(__eq__) |
| 3 | 1.717e-05 | 5.724e-06 | 0.2305 | 0.07685 | combine.py:52(update_fn) |
| 13 | 1.717e-05 | 1.321e-06 | 0.002476 | 0.0001905 | slicing.py:50(slice) |
| 203 | 1.716e-05 | 8.453e-08 | 1.716e-05 | 8.453e-08 | pxla.py:516(__repr__) |
| 13 | 1.708e-05 | 1.314e-06 | 0.002847 | 0.000219 | ad.py:552(<lambda>) |
| 1 | 1.701e-05 | 1.701e-05 | 3.828e-05 | 3.828e-05 | _stablehlo_ops_gen.py:4616(__init__) |
| 2 | 1.692e-05 | 8.459e-06 | 0.0001853 | 9.267e-05 | partial_eval.py:1098(<listcomp>) |
| 64 | 1.687e-05 | 2.636e-07 | 1.687e-05 | 2.636e-07 | partial_eval.py:975(<genexpr>) |
| 1 | 1.681e-05 | 1.681e-05 | 0.002705 | 0.002705 | lax_numpy.py:1182(clip) |
| 8 | 1.679e-05 | 2.099e-06 | 0.000805 | 0.0001006 | transform.py:70(to_prng_sequence) |
| 6 | 1.676e-05 | 2.794e-06 | 3.812e-05 | 6.353e-06 | prng.py:276(keys_aval_to_base_arr_aval) |
| 32 | 1.674e-05 | 5.23e-07 | 1.674e-05 | 5.23e-07 | initializers.py:94(__init__) |
| 15 | 1.672e-05 | 1.115e-06 | 1.672e-05 | 1.115e-06 | mlir.py:1018(<listcomp>) |
| 2 | 1.669e-05 | 8.345e-06 | 0.0002848 | 0.0001424 | pjit.py:1637(<listcomp>) |
| 8 | 1.638e-05 | 2.048e-06 | 6.956e-05 | 8.695e-06 | base.py:1041(extract_state) |
| 1 | 1.635e-05 | 1.635e-05 | 0.3988 | 0.3988 | algorithm.py:302(update_q) |
| 1 | 1.631e-05 | 1.631e-05 | 0.0004435 | 0.0004435 | passive_env_checker.py:170(env_reset_passive_checker) |
| 34 | 1.625e-05 | 4.778e-07 | 0.0001631 | 4.796e-06 | lax_numpy.py:313(issubdtype) |
| 16 | 1.599e-05 | 9.993e-07 | 0.001581 | 9.879e-05 | lax.py:364(sqrt) |
| 2 | 1.596e-05 | 7.978e-06 | 2.271e-05 | 1.136e-05 | partial_eval.py:1335(<listcomp>) |
| 12 | 1.581e-05 | 1.318e-06 | 0.002745 | 0.0002287 | ad.py:553(<lambda>) |
| 11 | 1.574e-05 | 1.431e-06 | 0.009456 | 0.0008597 | core.py:648(__rsub__) |
| 2 | 1.569e-05 | 7.844e-06 | 7.064e-05 | 3.532e-05 | _stablehlo_ops_gen.py:1566(__init__) |
| 8 | 1.552e-05 | 1.94e-06 | 0.0005316 | 6.645e-05 | core.py:732(__repr__) |
| 40 | 1.55e-05 | 3.875e-07 | 6.985e-05 | 1.746e-06 | mlir.py:1348(<genexpr>) |
| 12 | 1.496e-05 | 1.247e-06 | 1.861e-05 | 1.551e-06 | shape_base.py:23(atleast_1d) |
| 8 | 1.486e-05 | 1.858e-06 | 0.0003609 | 4.511e-05 | prng.py:931(_broadcast) |
| 22 | 1.466e-05 | 6.664e-07 | 0.001539 | 6.993e-05 | ad_util.py:36(add_jaxvals) |
| 3 | 1.464e-05 | 4.879e-06 | 0.004043 | 0.001348 | core.py:22(__init__) |
| 1 | 1.449e-05 | 1.449e-05 | 0.001296 | 0.001296 | util.py:429(_where) |
| 62 | 1.445e-05 | 2.331e-07 | 1.445e-05 | 2.331e-07 | mlir.py:913(<listcomp>) |
| 2 | 1.442e-05 | 7.21e-06 | 0.000158 | 7.902e-05 | partial_eval.py:1094(<listcomp>) |
| 13 | 1.44e-05 | 1.108e-06 | 4.076e-05 | 3.136e-06 | sharding_impls.py:614(__hash__) |
| 64 | 1.433e-05 | 2.239e-07 | 1.433e-05 | 2.239e-07 | pxla.py:2450(<listcomp>) |
| 26 | 1.429e-05 | 5.497e-07 | 6.544e-05 | 2.517e-06 | ad.py:123(<genexpr>) |
| 31 | 1.426e-05 | 4.601e-07 | 1.721e-05 | 5.552e-07 | partial_eval.py:1328(<listcomp>) |
| 46 | 1.419e-05 | 3.086e-07 | 2.291e-05 | 4.981e-07 | lax.py:1295(<genexpr>) |
| 4 | 1.4e-05 | 3.501e-06 | 2.14e-05 | 5.351e-06 | dtypes.py:100(_canonicalize_dtype) |
| 2 | 1.389e-05 | 6.946e-06 | 5.06e-05 | 2.53e-05 | api_util.py:254(argnums_partial) |
| 9 | 1.385e-05 | 1.539e-06 | 3.837e-05 | 4.263e-06 | prng.py:273(keys_shaped_array) |
| 8 | 1.383e-05 | 1.729e-06 | 6.836e-05 | 8.545e-06 | lax.py:1058(_get_max_identity) |
| 2 | 1.383e-05 | 6.916e-06 | 5.565e-05 | 2.782e-05 | partial_eval.py:1453(_move_binders_to_front) |
| 2 | 1.376e-05 | 6.882e-06 | 0.0002691 | 0.0001345 | sharding_impls.py:634(devices_indices_map) |
| 10 | 1.348e-05 | 1.348e-06 | 0.0003919 | 3.919e-05 | lax.py:3099(_squeeze_lower) |
| 1 | 1.344e-05 | 1.344e-05 | 0.004932 | 0.004932 | prng.py:1152(_threefry_random_bits_original) |
| 15 | 1.333e-05 | 8.885e-07 | 1.333e-05 | 8.885e-07 | pjit.py:1789(<listcomp>) |
| 16 | 1.332e-05 | 8.323e-07 | 0.0001256 | 7.85e-06 | reductions.py:101(<genexpr>) |
| 8 | 1.33e-05 | 1.663e-06 | 0.0001701 | 2.126e-05 | core.py:1449(__repr__) |
| 24 | 1.327e-05 | 5.529e-07 | 0.0002103 | 8.762e-06 | lax_numpy.py:743(isrealobj) |
| 54 | 1.322e-05 | 2.448e-07 | 1.322e-05 | 2.448e-07 | ~:0(<method 'splitlines' of 'str' objects>) |
| 9 | 1.321e-05 | 1.467e-06 | 1.321e-05 | 1.467e-06 | ~:0(<built-in method numpy.zeros>) |
| 4 | 1.311e-05 | 3.278e-06 | 1.667e-05 | 4.168e-06 | tree_util.py:365(__new__) |
| 2 | 1.293e-05 | 6.467e-06 | 4.097e-05 | 2.049e-05 | util.py:364(wrapper) |
| 12 | 1.288e-05 | 1.073e-06 | 4.282e-05 | 3.568e-06 | <__array_function__ internals>:177(copyto) |
| 12 | 1.261e-05 | 1.051e-06 | 0.0001762 | 1.468e-05 | lax.py:2691(<listcomp>) |
| 2 | 1.255e-05 | 6.275e-06 | 0.0003029 | 0.0001515 | lax.py:1209(broadcasted_iota) |
| 8 | 1.247e-05 | 1.559e-06 | 4.335e-05 | 5.418e-06 | source_info_util.py:260(transform_name_stack) |
| 62 | 1.242e-05 | 2.003e-07 | 1.242e-05 | 2.003e-07 | mlir.py:1202(<dictcomp>) |
| 4 | 1.226e-05 | 3.066e-06 | 0.0002571 | 6.427e-05 | lax.py:3117(_reshape_shape_rule) |
| 3 | 1.215e-05 | 4.051e-06 | 1.52e-05 | 5.067e-06 | copyreg.py:113(_slotnames) |
| 6 | 1.211e-05 | 2.019e-06 | 0.0001209 | 2.014e-05 | ad.py:707(f_jvp_traceable) |
| 4/2 | 1.194e-05 | 5.968e-06 | 0.03036 | 0.01518 | tree_util.py:302(__call__) |
| 6 | 1.19e-05 | 1.983e-06 | 0.2305 | 0.03842 | base.py:309(update) |
| 8 | 1.187e-05 | 1.484e-06 | 6.584e-05 | 8.23e-06 | lax.py:1068(_get_min_identity) |
| 2 | 1.179e-05 | 5.895e-06 | 3.4e-05 | 1.7e-05 | pxla.py:226(_op_sharding_to_numpy_indices) |
| 62 | 1.168e-05 | 1.884e-07 | 1.168e-05 | 1.884e-07 | mlir.py:914(<listcomp>) |
| 6 | 1.163e-05 | 1.938e-06 | 0.000152 | 2.533e-05 | pjit.py:1297(__eq__) |
| 11 | 1.161e-05 | 1.055e-06 | 0.002713 | 0.0002466 | lax.py:1466(<genexpr>) |
| 9 | 1.134e-05 | 1.259e-06 | 2.675e-05 | 2.972e-06 | partial_eval.py:1689(<lambda>) |
| 42 | 1.133e-05 | 2.696e-07 | 3.683e-05 | 8.77e-07 | lax.py:1555(<genexpr>) |
| 1 | 1.125e-05 | 1.125e-05 | 0.009744 | 0.009744 | mlir.py:97(delegate_lowering) |
| 96 | 1.125e-05 | 1.172e-07 | 1.125e-05 | 1.172e-07 | inspect.py:3035(return_annotation) |
| 4 | 1.115e-05 | 2.787e-06 | 0.0001378 | 3.445e-05 | partial_eval.py:160(new_instantiated_literal) |
| 4 | 1.1e-05 | 2.749e-06 | 0.0001518 | 3.795e-05 | core.py:1969(divide_shape_sizes) |
| 3 | 1.087e-05 | 3.622e-06 | 0.0001993 | 6.642e-05 | asserts_internal.py:231(_chex_assert_fn) |
| 62 | 1.084e-05 | 1.749e-07 | 1.084e-05 | 1.749e-07 | mlir.py:1205(<listcomp>) |
| 2 | 1.079e-05 | 5.393e-06 | 6.892e-05 | 3.446e-05 | partial_eval.py:1100(<listcomp>) |
| 15 | 1.069e-05 | 7.125e-07 | 1.069e-05 | 7.125e-07 | mlir.py:1056(<listcomp>) |
| 2 | 1.067e-05 | 5.334e-06 | 1.067e-05 | 5.334e-06 | partial_eval.py:1334(<dictcomp>) |
| 1 | 1.065e-05 | 1.065e-05 | 0.009547 | 0.009547 | random.py:243(uniform) |
| 4 | 1.059e-05 | 2.647e-06 | 1.337e-05 | 3.343e-06 | util.py:193(<listcomp>) |
| 62 | 1.053e-05 | 1.698e-07 | 1.053e-05 | 1.698e-07 | mlir.py:821(<listcomp>) |
| 1 | 1.045e-05 | 1.045e-05 | 0.1691 | 0.1691 | algorithm.py:289(update_pi) |
| 1 | 1.041e-05 | 1.041e-05 | 0.00965 | 0.00965 | prng.py:699(random_bits_lowering) |
| 2 | 1.037e-05 | 5.186e-06 | 8.616e-05 | 4.308e-05 | passive_env_checker.py:122(check_obs) |
| 2 | 1.033e-05 | 5.165e-06 | 1.36e-05 | 6.802e-06 | ad.py:145(<listcomp>) |
| 8 | 1.012e-05 | 1.264e-06 | 3.679e-05 | 4.598e-06 | __init__.py:321(filter) |
| 2 | 1.01e-05 | 5.051e-06 | 0.01999 | 0.009994 | transform.py:110(bias_correction) |
| 1 | 1.006e-05 | 1.006e-05 | 1.827e-05 | 1.827e-05 | shape_base.py:1191(tile) |
| 1 | 1.001e-05 | 1.001e-05 | 0.004936 | 0.004936 | prng.py:1099(_threefry_split_original) |
| 46 | 9.89e-06 | 2.15e-07 | 9.89e-06 | 2.15e-07 | lax.py:3064(_squeeze_dtype_rule) |
| 4 | 9.86e-06 | 2.465e-06 | 7.619e-05 | 1.905e-05 | lax.py:2313(_bitcast_convert_type_dtype_rule) |
| 4 | 9.793e-06 | 2.448e-06 | 1.614e-05 | 4.036e-06 | source_info_util.py:84(transform) |
| 85 | 9.687e-06 | 1.14e-07 | 9.687e-06 | 1.14e-07 | dispatch.py:339(<genexpr>) |
| 62 | 9.632e-06 | 1.554e-07 | 9.632e-06 | 1.554e-07 | mlir.py:1203(<dictcomp>) |
| 2 | 9.612e-06 | 4.806e-06 | 4.294e-05 | 2.147e-05 | _stablehlo_ops_gen.py:2755(__init__) |
| 4 | 9.6e-06 | 2.4e-06 | 0.0001545 | 3.864e-05 | ad.py:298(pure) |
| 2 | 9.573e-06 | 4.787e-06 | 0.0001231 | 6.156e-05 | partial_eval.py:1096(<listcomp>) |
| 4 | 9.478e-06 | 2.37e-06 | 5.387e-05 | 1.347e-05 | prng.py:683(random_bits_abstract_eval) |
| 1 | 9.451e-06 | 9.451e-06 | 0.007281 | 0.007281 | prng.py:1121(threefry_random_bits) |
| 2 | 9.429e-06 | 4.715e-06 | 0.001032 | 0.0005162 | lax_numpy.py:1158(split) |
| 4 | 9.409e-06 | 2.352e-06 | 0.0001085 | 2.712e-05 | core.py:1847(divide_shape_sizes) |
| 96 | 9.404e-06 | 9.796e-08 | 9.404e-06 | 9.796e-08 | contextlib.py:543(__enter__) |
| 1 | 9.387e-06 | 9.387e-06 | 0.02134 | 0.02134 | algorithm.py:266(compute_loss_pi) |
| 12 | 9.375e-06 | 7.813e-07 | 3.809e-05 | 3.175e-06 | <__array_function__ internals>:177(atleast_1d) |
| 2 | 9.285e-06 | 4.643e-06 | 3.443e-05 | 1.722e-05 | pjit.py:1563(_known_jaxpr_fwd) |
| 26 | 9.24e-06 | 3.554e-07 | 9.24e-06 | 3.554e-07 | partial_eval.py:1068(<genexpr>) |
| 3 | 9.232e-06 | 3.077e-06 | 0.05739 | 0.01913 | transform.py:81(update_moment) |
| 2 | 9.12e-06 | 4.56e-06 | 2.632e-05 | 1.316e-05 | traceback_util.py:139(api_boundary) |
| 2 | 9.058e-06 | 4.529e-06 | 9.058e-06 | 4.529e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.get_tuple>) |
| 2 | 9.028e-06 | 4.514e-06 | 1.676e-05 | 8.379e-06 | _ufunc_config.py:33(seterr) |
| 2 | 9.014e-06 | 4.507e-06 | 0.0002716 | 0.0001358 | lax.py:2049(_sub_jvp) |
| 4 | 8.928e-06 | 2.232e-06 | 0.0004592 | 0.0001148 | lax.py:1084(_reduce_max) |
| 55 | 8.895e-06 | 1.617e-07 | 8.895e-06 | 1.617e-07 | function_base.py:5050(_delete_dispatcher) |
| 2 | 8.889e-06 | 4.445e-06 | 0.0001521 | 7.603e-05 | lax_numpy.py:758(ravel) |
| 6 | 8.816e-06 | 1.469e-06 | 5.199e-05 | 8.665e-06 | lax.py:1619(_maybe_broadcast) |
| 16 | 8.809e-06 | 5.506e-07 | 0.000751 | 4.694e-05 | data_structures.py:155(to_dict) |
| 8 | 8.744e-06 | 1.093e-06 | 1.665e-05 | 2.081e-06 | pretty_printer.py:353(text) |
| 6 | 8.735e-06 | 1.456e-06 | 3.959e-05 | 6.598e-06 | prng.py:740(random_wrap_abstract_eval) |
| 2 | 8.699e-06 | 4.35e-06 | 0.0001157 | 5.785e-05 | partial_eval.py:1099(<listcomp>) |
| 30 | 8.693e-06 | 2.898e-07 | 1.038e-05 | 3.459e-07 | prng.py:841(<genexpr>) |
| 2 | 8.564e-06 | 4.282e-06 | 0.0001113 | 5.566e-05 | lax_numpy.py:748(reshape) |
| 24 | 8.519e-06 | 3.55e-07 | 1.123e-05 | 4.679e-07 | ad.py:280(<lambda>) |
| 3 | 8.509e-06 | 2.836e-06 | 0.04515 | 0.01505 | transform.py:93(update_moment_per_elem_norm) |
| 2 | 8.412e-06 | 4.206e-06 | 0.0006746 | 0.0003373 | lax_numpy.py:1155(<listcomp>) |
| 16 | 8.38e-06 | 5.238e-07 | 6.402e-05 | 4.001e-06 | transform.py:382(check_mapping) |
| 3 | 8.325e-06 | 2.775e-06 | 0.0004695 | 0.0001565 | lax.py:2104(_div_transpose_rule) |
| 35 | 8.266e-06 | 2.362e-07 | 0.0002597 | 7.42e-06 | ad.py:581(instantiate_zeros) |
| 38 | 8.239e-06 | 2.168e-07 | 8.239e-06 | 2.168e-07 | ~:0(<method 'getvalue' of '_io.StringIO' objects>) |
| 94 | 8.22e-06 | 8.745e-08 | 8.22e-06 | 8.745e-08 | pjit.py:1578(<genexpr>) |
| 48 | 8.199e-06 | 1.708e-07 | 8.199e-06 | 1.708e-07 | utils.py:82(<genexpr>) |
| 2 | 8.186e-06 | 4.093e-06 | 8.186e-06 | 4.093e-06 | pjit.py:1536(<listcomp>) |
| 16 | 8.069e-06 | 5.043e-07 | 6.557e-05 | 4.098e-06 | lax.py:4335(<genexpr>) |
| 16 | 8.034e-06 | 5.021e-07 | 0.0007591 | 4.744e-05 | data_structures.py:127(to_haiku_dict) |
| 1 | 7.94e-06 | 7.94e-06 | 0.009767 | 0.009767 | prng.py:612(random_split_lowering) |
| 7 | 7.913e-06 | 1.13e-06 | 3.632e-05 | 5.188e-06 | partial_eval.py:1370(has_effects) |
| 2 | 7.897e-06 | 3.949e-06 | 1e-05 | 5.002e-06 | pjit.py:1626(<listcomp>) |
| 4 | 7.862e-06 | 1.966e-06 | 0.0004526 | 0.0001131 | lax.py:1087(_reduce_min) |
| 36 | 7.84e-06 | 2.178e-07 | 7.84e-06 | 2.178e-07 | lax.py:3226(<genexpr>) |
| 11 | 7.838e-06 | 7.125e-07 | 0.0005714 | 5.194e-05 | util.py:238(<listcomp>) |
| 15 | 7.799e-06 | 5.199e-07 | 7.799e-06 | 5.199e-07 | pjit.py:1790(<listcomp>) |
| 3 | 7.769e-06 | 2.59e-06 | 0.01636 | 0.005453 | update.py:24(apply_updates) |
| 8 | 7.474e-06 | 9.343e-07 | 7.703e-05 | 9.629e-06 | base.py:189(collect_state) |
| 50 | 7.443e-06 | 1.489e-07 | 7.443e-06 | 1.489e-07 | fromnumeric.py:1034(_argsort_dispatcher) |
| 12 | 7.399e-06 | 6.166e-07 | 7.399e-06 | 6.166e-07 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.isinstance>) |
| 3 | 7.308e-06 | 2.436e-06 | 0.0004767 | 0.0001589 | lax.py:3082(_squeeze_transpose_rule) |
| 2 | 7.286e-06 | 3.643e-06 | 7.286e-06 | 3.643e-06 | partial_eval.py:1081(<listcomp>) |
| 11 | 7.26e-06 | 6.6e-07 | 1.473e-05 | 1.339e-06 | core.py:307(__init__) |
| 2 | 7.257e-06 | 3.629e-06 | 7.257e-06 | 3.629e-06 | partial_eval.py:1092(<listcomp>) |
| 4 | 7.16e-06 | 1.79e-06 | 1.37e-05 | 3.424e-06 | partial_eval.py:805(<listcomp>) |
| 2 | 7.062e-06 | 3.531e-06 | 1.388e-05 | 6.942e-06 | sharding_impls.py:602(_op_sharding_hash) |
| 8 | 7.036e-06 | 8.795e-07 | 0.0007332 | 9.166e-05 | lax.py:1395(square) |
| 2 | 7.026e-06 | 3.513e-06 | 8.714e-05 | 4.357e-05 | array_methods.py:131(_reshape) |
| 2 | 7.016e-06 | 3.508e-06 | 2.995e-05 | 1.498e-05 | hlo_helpers.py:60(<listcomp>) |
| 3 | 6.98e-06 | 2.327e-06 | 0.0139 | 0.004634 | transform.py:516(update_fn) |
| 3 | 6.972e-06 | 2.324e-06 | 0.0001884 | 6.28e-05 | asserts_internal.py:159(_assert_on_host) |
| 4 | 6.904e-06 | 1.726e-06 | 1.854e-05 | 4.636e-06 | lax.py:2295(_bitcast_convert_type_shape_rule) |
| 1 | 6.836e-06 | 6.836e-06 | 9.122e-06 | 9.122e-06 | dtypes.py:349(__new__) |
| 8 | 6.814e-06 | 8.518e-07 | 7.903e-06 | 9.879e-07 | pretty_printer.py:84(__init__) |
| 4 | 6.799e-06 | 1.7e-06 | 0.0003212 | 8.029e-05 | lax.py:4353(_iota_staging_rule) |
| 8 | 6.787e-06 | 8.484e-07 | 2.545e-05 | 3.181e-06 | pjit.py:1540(_filter_zeros) |
| 4 | 6.71e-06 | 1.678e-06 | 1.945e-05 | 4.863e-06 | ad.py:62(jvp) |
| 4 | 6.621e-06 | 1.655e-06 | 0.000435 | 0.0001088 | lax.py:3175(_reshape_staging_rule) |
| 38 | 6.594e-06 | 1.735e-07 | 6.594e-06 | 1.735e-07 | ~:0(<method 'strip' of 'str' objects>) |
| 2 | 6.562e-06 | 3.281e-06 | 3.601e-05 | 1.801e-05 | array_methods.py:115(_compute_newshape) |
| 1 | 6.49e-06 | 6.49e-06 | 0.001012 | 0.001012 | core.py:636(__neg__) |
| 2 | 6.382e-06 | 3.191e-06 | 0.0007875 | 0.0003938 | lax.py:1949(_integer_pow_jvp) |
| 8 | 6.319e-06 | 7.899e-07 | 2.14e-05 | 2.675e-06 | core.py:490(__repr__) |
| 2 | 6.177e-06 | 3.089e-06 | 1.042e-05 | 5.209e-06 | __init__.py:272(__getattr__) |
| 37 | 6.165e-06 | 1.666e-07 | 6.165e-06 | 1.666e-07 | pjit.py:1614(<genexpr>) |
| 16 | 6.164e-06 | 3.853e-07 | 1.51e-05 | 9.437e-07 | partial_eval.py:798(<genexpr>) |
| 38 | 6.142e-06 | 1.616e-07 | 6.142e-06 | 1.616e-07 | ~:0(<method 'write' of '_io.StringIO' objects>) |
| 2 | 6.125e-06 | 3.063e-06 | 2.724e-05 | 1.362e-05 | partial_eval.py:1464(_move_to_front) |
| 1 | 6.114e-06 | 6.114e-06 | 3.968e-05 | 3.968e-05 | lax.py:2333(_bitcast_convert_type_lower) |
| 23 | 6.04e-06 | 2.626e-07 | 6.04e-06 | 2.626e-07 | core.py:897(__init__) |
| 16 | 5.988e-06 | 3.743e-07 | 5.988e-06 | 3.743e-07 | pretty_printer.py:258(<genexpr>) |
| 38 | 5.978e-06 | 1.573e-07 | 5.978e-06 | 1.573e-07 | ~:0(<method 'copy' of 'dict' objects>) |
| 2 | 5.961e-06 | 2.981e-06 | 0.02431 | 0.01216 | ad.py:262(closed_backward_pass) |
| 2 | 5.938e-06 | 2.969e-06 | 5.938e-06 | 2.969e-06 | pjit.py:1575(<listcomp>) |
| 2 | 5.838e-06 | 2.919e-06 | 5.838e-06 | 2.919e-06 | ~:0(<method 'fill' of 'numpy.ndarray' objects>) |
| 1 | 5.786e-06 | 5.786e-06 | 8.271e-06 | 8.271e-06 | random.py:800(getrandbits) |
| 1 | 5.696e-06 | 5.696e-06 | 1.148e-05 | 1.148e-05 | inspect.py:2037(_signature_bound_method) |
| 30 | 5.461e-06 | 1.82e-07 | 5.461e-06 | 1.82e-07 | ~:0(<method 'keys' of 'collections.OrderedDict' objects>) |
| 27 | 5.411e-06 | 2.004e-07 | 5.411e-06 | 2.004e-07 | util.py:290(__eq__) |
| 2 | 5.402e-06 | 2.701e-06 | 5.402e-06 | 2.701e-06 | partial_eval.py:1465(<listcomp>) |
| 8 | 5.383e-06 | 6.729e-07 | 1.038e-05 | 1.298e-06 | lax_numpy.py:1154(<lambda>) |
| 2 | 5.348e-06 | 2.674e-06 | 5.348e-06 | 2.674e-06 | ad.py:702(<listcomp>) |
| 2 | 5.296e-06 | 2.648e-06 | 4.304e-05 | 2.152e-05 | lax.py:4586(_const) |
| 12 | 5.294e-06 | 4.412e-07 | 6.968e-05 | 5.807e-06 | lax.py:3118(<genexpr>) |
| 2 | 5.254e-06 | 2.627e-06 | 2.178e-05 | 1.089e-05 | pjit.py:1658(<listcomp>) |
| 2 | 5.143e-06 | 2.572e-06 | 5.143e-06 | 2.572e-06 | partial_eval.py:1466(<listcomp>) |
| 2 | 5.117e-06 | 2.559e-06 | 5.984e-05 | 2.992e-05 | api.py:2267(<listcomp>) |
| 28 | 5.076e-06 | 1.813e-07 | 5.076e-06 | 1.813e-07 | multiarray.py:345(where) |
| 2 | 5.05e-06 | 2.525e-06 | 0.0002753 | 0.0001376 | lax.py:788(broadcast_to_rank) |
| 1 | 5.037e-06 | 5.037e-06 | 5.037e-06 | 5.037e-06 | ~:0(<method 'cumsum' of 'numpy.ndarray' objects>) |
| 3 | 5.032e-06 | 1.677e-06 | 5.509e-05 | 1.836e-05 | core.py:597(__iter__) |
| 1 | 5.032e-06 | 5.032e-06 | 1.733e-05 | 1.733e-05 | _ufunc_config.py:430(__enter__) |
| 8 | 5e-06 | 6.25e-07 | 5e-06 | 6.25e-07 | util.py:106(subvals) |
| 2 | 4.967e-06 | 2.484e-06 | 2.547e-05 | 1.274e-05 | api.py:785(_check_output_dtype_revderiv) |
| 17 | 4.894e-06 | 2.879e-07 | 4.894e-06 | 2.879e-07 | ad.py:498(<genexpr>) |
| 12 | 4.861e-06 | 4.051e-07 | 6.536e-06 | 5.447e-07 | partial_eval.py:714(is_known) |
| 15 | 4.806e-06 | 3.204e-07 | 4.806e-06 | 3.204e-07 | pjit.py:1457(<listcomp>) |
| 27 | 4.802e-06 | 1.779e-07 | 4.802e-06 | 1.779e-07 | fromnumeric.py:3190(_size_dispatcher) |
| 1 | 4.785e-06 | 4.785e-06 | 8.399e-05 | 8.399e-05 | lax.py:1268(stop) |
| 2 | 4.733e-06 | 2.367e-06 | 0.02693 | 0.01347 | pjit.py:1715(_pjit_transpose_trace) |
| 16 | 4.676e-06 | 2.923e-07 | 4.676e-06 | 2.923e-07 | pretty_printer.py:240(_update_color) |
| 16 | 4.622e-06 | 2.889e-07 | 4.622e-06 | 2.889e-07 | pretty_printer.py:337(<genexpr>) |
| 16 | 4.591e-06 | 2.869e-07 | 4.591e-06 | 2.869e-07 | ~:0(<method 'lower' of 'str' objects>) |
| 2 | 4.55e-06 | 2.275e-06 | 4.829e-06 | 2.415e-06 | _ufunc_config.py:132(geterr) |
| 2 | 4.462e-06 | 2.231e-06 | 4.462e-06 | 2.231e-06 | pjit.py:1567(<listcomp>) |
| 40 | 4.451e-06 | 1.113e-07 | 4.451e-06 | 1.113e-07 | __init__.py:599(__missing__) |
| 6 | 4.43e-06 | 7.383e-07 | 4.43e-06 | 7.383e-07 | ~:0(<method 'intersection' of 'set' objects>) |
| 10 | 4.425e-06 | 4.425e-07 | 4.425e-06 | 4.425e-07 | api_util.py:162(__eq__) |
| 17 | 4.422e-06 | 2.601e-07 | 4.422e-06 | 2.601e-07 | dtypes.py:540(<genexpr>) |
| 3 | 4.403e-06 | 1.468e-06 | 9.003e-05 | 3.001e-05 | lax_numpy.py:325(result_type) |
| 1 | 4.39e-06 | 4.39e-06 | 0.000499 | 0.000499 | ufuncs.py:307(_power) |
| 30 | 4.385e-06 | 1.462e-07 | 4.385e-06 | 1.462e-07 | prng.py:843(<genexpr>) |
| 8 | 4.368e-06 | 5.46e-07 | 4.368e-06 | 5.46e-07 | core.py:900(__repr__) |
| 2 | 4.27e-06 | 2.135e-06 | 1.208e-05 | 6.04e-06 | api_util.py:456(check_callable) |
| 2 | 4.263e-06 | 2.132e-06 | 4.263e-06 | 2.132e-06 | lax_numpy.py:1144(<listcomp>) |
| 2 | 4.255e-06 | 2.127e-06 | 1.307e-05 | 6.533e-06 | api_util.py:60(_ensure_index) |
| 2 | 4.24e-06 | 2.12e-06 | 4.24e-06 | 2.12e-06 | ~:0(<built-in method _warnings.warn>) |
| 2 | 4.151e-06 | 2.076e-06 | 4.151e-06 | 2.076e-06 | ~:0(<built-in method jaxlib.cuda._prng.threefry2x32_descriptor>) |
| 2 | 4.099e-06 | 2.05e-06 | 6.316e-06 | 3.158e-06 | api_util.py:446(_isgeneratorfunction) |
| 3 | 4.051e-06 | 1.35e-06 | 2.405e-05 | 8.017e-06 | prng.py:771(random_unwrap_abstract_eval) |
| 19 | 4.049e-06 | 2.131e-07 | 4.049e-06 | 2.131e-07 | tree_util.py:238(<lambda>) |
| 2 | 4.047e-06 | 2.024e-06 | 0.000111 | 5.551e-05 | partial_eval.py:1468(move_binders_to_back) |
| 2 | 4.019e-06 | 2.009e-06 | 7.383e-05 | 3.691e-05 | space.py:112(__contains__) |
| 1 | 3.968e-06 | 3.968e-06 | 8.262e-06 | 8.262e-06 | prng.py:788(_is_threefry_prng_key) |
| 1 | 3.963e-06 | 3.963e-06 | 1.912e-05 | 1.912e-05 | _stablehlo_ops_gen.py:675(__init__) |
| 1 | 3.959e-06 | 3.959e-06 | 2.535e-05 | 2.535e-05 | lax.py:3261(_select_dtype_rule) |
| 8 | 3.945e-06 | 4.931e-07 | 1.357e-05 | 1.696e-06 | reductions.py:45(_asarray) |
| 1 | 3.942e-06 | 3.942e-06 | 8.089e-05 | 8.089e-05 | prng.py:664(random_bits) |
| 2 | 3.883e-06 | 1.941e-06 | 3.883e-06 | 1.941e-06 | partial_eval.py:1091(<listcomp>) |
| 27 | 3.816e-06 | 1.413e-07 | 3.816e-06 | 1.413e-07 | <frozen _collections_abc>:409(__subclasshook__) |
| 25 | 3.793e-06 | 1.517e-07 | 3.793e-06 | 1.517e-07 | fromnumeric.py:89(_take_dispatcher) |
| 1 | 3.666e-06 | 3.666e-06 | 0.000354 | 0.000354 | util.py:375(promote_args_numeric) |
| 1 | 3.642e-06 | 3.642e-06 | 0.0005058 | 0.0005058 | ufuncs.py:107(fn) |
| 2 | 3.613e-06 | 1.807e-06 | 0.0003065 | 0.0001533 | lax.py:1202(iota) |
| 3 | 3.565e-06 | 1.188e-06 | 0.004337 | 0.001446 | core.py:649(__mul__) |
| 3 | 3.521e-06 | 1.174e-06 | 4.633e-06 | 1.544e-06 | core.py:2134(__hash__) |
| 4 | 3.518e-06 | 8.795e-07 | 0.00371 | 0.0009276 | core.py:647(__sub__) |
| 3 | 3.506e-06 | 1.169e-06 | 1.737e-05 | 5.79e-06 | api_util.py:138(flatten_fun_nokwargs2) |
| 4 | 3.505e-06 | 8.763e-07 | 0.0001553 | 3.883e-05 | core.py:1977(same_shape_sizes) |
| 2 | 3.4e-06 | 1.7e-06 | 4.531e-06 | 2.266e-06 | partial_eval.py:1075(<listcomp>) |
| 3 | 3.39e-06 | 1.13e-06 | 0.00282 | 0.00094 | core.py:640(__lt__) |
| 4 | 3.284e-06 | 8.21e-07 | 3.284e-06 | 8.21e-07 | partial_eval.py:803(<listcomp>) |
| 8 | 3.257e-06 | 4.071e-07 | 3.981e-06 | 4.976e-07 | lax.py:4343(<genexpr>) |
| 2 | 3.224e-06 | 1.612e-06 | 6.648e-06 | 3.324e-06 | api.py:750(_check_scalar) |
| 2 | 3.22e-06 | 1.61e-06 | 3.894e-06 | 1.947e-06 | pxla.py:3156(<listcomp>) |
| 4 | 3.2e-06 | 8e-07 | 3.77e-06 | 9.425e-07 | ad.py:82(<listcomp>) |
| 2 | 3.116e-06 | 1.558e-06 | 4.826e-06 | 2.413e-06 | partial_eval.py:1078(<listcomp>) |
| 1 | 3.111e-06 | 3.111e-06 | 0.000193 | 0.000193 | util.py:390(_broadcast_arrays) |
| 3 | 2.995e-06 | 9.983e-07 | 1.603e-05 | 5.345e-06 | prng.py:598(random_split_abstract_eval) |
| 1 | 2.981e-06 | 2.981e-06 | 0.00953 | 0.00953 | core.py:112(__call__) |
| 1 | 2.977e-06 | 2.977e-06 | 8.532e-05 | 8.532e-05 | random.py:92(_random_bits) |
| 15 | 2.924e-06 | 1.949e-07 | 2.924e-06 | 1.949e-07 | mlir.py:920(<listcomp>) |
| 2 | 2.903e-06 | 1.451e-06 | 7.391e-06 | 3.696e-06 | api_util.py:81(_ensure_str_tuple) |
| 24 | 2.875e-06 | 1.198e-07 | 2.875e-06 | 1.198e-07 | ad.py:279(<lambda>) |
| 1 | 2.846e-06 | 2.846e-06 | 0.0001992 | 0.0001992 | lax.py:2946(<listcomp>) |
| 2 | 2.833e-06 | 1.417e-06 | 8.069e-05 | 4.035e-05 | lax.py:3167(_reshape_lower) |
| 1 | 2.826e-06 | 2.826e-06 | 6.035e-06 | 6.035e-06 | tree_util.py:595(<lambda>) |
| 2 | 2.824e-06 | 1.412e-06 | 8.786e-06 | 4.393e-06 | ad.py:317(<listcomp>) |
| 1 | 2.815e-06 | 2.815e-06 | 5.552e-06 | 5.552e-06 | lax.py:2930(<listcomp>) |
| 2 | 2.794e-06 | 1.397e-06 | 4.529e-06 | 2.265e-06 | pjit.py:1605(<listcomp>) |
| 7 | 2.726e-06 | 3.894e-07 | 7.951e-06 | 1.136e-06 | lax.py:1462(<genexpr>) |
| 1 | 2.725e-06 | 2.725e-06 | 0.0002358 | 0.0002358 | ufuncs.py:62(<lambda>) |
| 1 | 2.723e-06 | 2.723e-06 | 0.0002514 | 0.0002514 | ad_util.py:57(zeros_like_aval) |
| 3 | 2.705e-06 | 9.017e-07 | 5.805e-06 | 1.935e-06 | core.py:2148(as_named_shape) |
| 2 | 2.681e-06 | 1.341e-06 | 7.659e-06 | 3.83e-06 | <__array_function__ internals>:177(can_cast) |
| 1 | 2.675e-06 | 2.675e-06 | 0.0002487 | 0.0002487 | lax.py:1192(zeros_like_shaped_array) |
| 2 | 2.634e-06 | 1.317e-06 | 0.000325 | 0.0001625 | lax.py:434(max) |
| 12 | 2.501e-06 | 2.084e-07 | 2.501e-06 | 2.084e-07 | multiarray.py:1079(copyto) |
| 2 | 2.494e-06 | 1.247e-06 | 1.167e-05 | 5.835e-06 | partial_eval.py:1095(<listcomp>) |
| 4 | 2.464e-06 | 6.16e-07 | 2.464e-06 | 6.16e-07 | partial_eval.py:785(<listcomp>) |
| 2 | 2.413e-06 | 1.207e-06 | 2.996e-06 | 1.498e-06 | dtypes.py:630(<setcomp>) |
| 6 | 2.412e-06 | 4.02e-07 | 2.412e-06 | 4.02e-07 | lax.py:2833(_clamp_shape_rule) |
| 9 | 2.384e-06 | 2.649e-07 | 2.384e-06 | 2.649e-07 | prng.py:456(__init__) |
| 2 | 2.332e-06 | 1.166e-06 | 1.159e-05 | 5.794e-06 | partial_eval.py:1104(<listcomp>) |
| 2 | 2.312e-06 | 1.156e-06 | 2.312e-06 | 1.156e-06 | ~:0(<built-in method jaxlib.mlir._mlir_libs._mlir.ir.get_unsigned>) |
| 8 | 2.271e-06 | 2.839e-07 | 2.271e-06 | 2.839e-07 | ad.py:183(<genexpr>) |
| 1 | 2.208e-06 | 2.208e-06 | 0.007283 | 0.007283 | prng.py:703(<lambda>) |
| 10 | 2.206e-06 | 2.206e-07 | 2.206e-06 | 2.206e-07 | pjit.py:1582(<genexpr>) |
| 8 | 2.134e-06 | 2.668e-07 | 2.134e-06 | 2.668e-07 | core.py:722(<listcomp>) |
| 8 | 2.122e-06 | 2.653e-07 | 2.122e-06 | 2.653e-07 | ~:0(<method 'group' of 're.Match' objects>) |
| 2 | 2.12e-06 | 1.06e-06 | 1.657e-05 | 8.287e-06 | mlir.py:514(replace) |
| 3 | 2.108e-06 | 7.027e-07 | 0.0003359 | 0.000112 | lax.py:2112(<lambda>) |
| 4 | 2.106e-06 | 5.265e-07 | 2.844e-05 | 7.111e-06 | partial_eval.py:146(pure) |
| 2 | 2.102e-06 | 1.051e-06 | 2.458e-06 | 1.229e-06 | pjit.py:1558(<listcomp>) |
| 1 | 2.101e-06 | 2.101e-06 | 2.101e-06 | 2.101e-06 | ~:0(<built-in method posix.urandom>) |
| 2 | 2.101e-06 | 1.051e-06 | 4.765e-06 | 2.383e-06 | core.py:207(replace) |
| 8 | 2.091e-06 | 2.614e-07 | 2.091e-06 | 2.614e-07 | base.py:256(<dictcomp>) |
| 2 | 2.031e-06 | 1.016e-06 | 0.000285 | 0.0001425 | lax.py:441(min) |
| 1 | 2.015e-06 | 2.015e-06 | 1.729e-05 | 1.729e-05 | util.py:393(<listcomp>) |
| 4 | 2.008e-06 | 5.02e-07 | 2.4e-06 | 6e-07 | ad.py:130(<genexpr>) |
| 8 | 1.995e-06 | 2.494e-07 | 1.995e-06 | 2.494e-07 | base.py:1042(<dictcomp>) |
| 2 | 1.976e-06 | 9.88e-07 | 1.976e-06 | 9.88e-07 | ~:0(<built-in method numpy.seterrobj>) |
| 3 | 1.974e-06 | 6.58e-07 | 1.974e-06 | 6.58e-07 | ad.py:99(<listcomp>) |
| 1 | 1.972e-06 | 1.972e-06 | 4.656e-06 | 4.656e-06 | ad.py:111(<listcomp>) |
| 4 | 1.971e-06 | 4.928e-07 | 5.223e-06 | 1.306e-06 | util.py:394(<genexpr>) |
| 2 | 1.961e-06 | 9.805e-07 | 3.018e-06 | 1.509e-06 | pxla.py:3250(<listcomp>) |
| 4 | 1.944e-06 | 4.86e-07 | 2.377e-06 | 5.943e-07 | ad.py:144(<genexpr>) |
| 4 | 1.942e-06 | 4.855e-07 | 3.088e-06 | 7.72e-07 | tree_util.py:595(<genexpr>) |
| 2 | 1.916e-06 | 9.58e-07 | 5.011e-06 | 2.505e-06 | lax.py:1213(<listcomp>) |
| 4 | 1.899e-06 | 4.748e-07 | 1.15e-05 | 2.874e-06 | lax.py:1487(<lambda>) |
| 1 | 1.873e-06 | 1.873e-06 | 4.029e-06 | 4.029e-06 | lax.py:3249(_select_shape_rule) |
| 1 | 1.86e-06 | 1.86e-06 | 9.538e-05 | 9.538e-05 | lax.py:1248(stop_gradient) |
| 2 | 1.803e-06 | 9.015e-07 | 0.0001799 | 8.996e-05 | lax.py:207(neg) |
| 9 | 1.802e-06 | 2.002e-07 | 1.802e-06 | 2.002e-07 | ad.py:579(<lambda>) |
| 2 | 1.795e-06 | 8.975e-07 | 1.212e-05 | 6.06e-06 | api.py:2268(<listcomp>) |
| 2 | 1.754e-06 | 8.77e-07 | 0.0004625 | 0.0002312 | lax.py:569(clamp) |
| 2 | 1.748e-06 | 8.74e-07 | 1.748e-06 | 8.74e-07 | util.py:354(wraps) |
| 4 | 1.743e-06 | 4.358e-07 | 1.743e-06 | 4.358e-07 | contextlib.py:751(__init__) |
| 12 | 1.735e-06 | 1.446e-07 | 1.735e-06 | 1.446e-07 | shape_base.py:19(_atleast_1d_dispatcher) |
| 4 | 1.729e-06 | 4.323e-07 | 1.729e-06 | 4.323e-07 | ~:0(<built-in method builtins.divmod>) |
| 2 | 1.691e-06 | 8.455e-07 | 3.214e-06 | 1.607e-06 | lax.py:172(_extract_tracers_dyn_shape) |
| 2 | 1.69e-06 | 8.45e-07 | 1.088e-05 | 5.442e-06 | partial_eval.py:1103(<listcomp>) |
| 1 | 1.671e-06 | 1.671e-06 | 1.984e-06 | 1.984e-06 | getlimits.py:476(__new__) |
| 1 | 1.661e-06 | 1.661e-06 | 1.661e-06 | 1.661e-06 | ~:0(<method 'repeat' of 'numpy.ndarray' objects>) |
| 3 | 1.66e-06 | 5.533e-07 | 5.514e-06 | 1.838e-06 | prng.py:469(__hash__) |
| 3 | 1.654e-06 | 5.513e-07 | 1.654e-06 | 5.513e-07 | ~:0(<method 'tolist' of 'numpy.ndarray' objects>) |
| 1 | 1.633e-06 | 1.633e-06 | 1.633e-06 | 1.633e-06 | ~:0(<method 'view' of 'numpy.ndarray' objects>) |
| 4 | 1.625e-06 | 4.063e-07 | 2.58e-05 | 6.451e-06 | array_methods.py:126(<genexpr>) |
| 1 | 1.609e-06 | 1.609e-06 | 0.0001649 | 0.0001649 | util.py:395(<listcomp>) |
| 1 | 1.603e-06 | 1.603e-06 | 3.112e-06 | 3.112e-06 | lax.py:3271(_select_weak_type_rule) |
| 2 | 1.554e-06 | 7.77e-07 | 0.0002241 | 0.0001121 | lax.py:306(tanh) |
| 4 | 1.54e-06 | 3.85e-07 | 1.54e-06 | 3.85e-07 | ~:0(<method 'get' of 'mappingproxy' objects>) |
| 1 | 1.523e-06 | 1.523e-06 | 5.983e-06 | 5.983e-06 | _ufunc_config.py:435(__exit__) |
| 4 | 1.491e-06 | 3.727e-07 | 1.491e-06 | 3.727e-07 | tree_util.py:299(__init__) |
| 1 | 1.491e-06 | 1.491e-06 | 4.851e-06 | 4.851e-06 | inspect.py:3039(replace) |
| 5 | 1.48e-06 | 2.96e-07 | 1.48e-06 | 2.96e-07 | api_util.py:270(<genexpr>) |
| 4 | 1.43e-06 | 3.575e-07 | 2.45e-06 | 6.125e-07 | pxla.py:2488(<genexpr>) |
| 6 | 1.427e-06 | 2.378e-07 | 1.427e-06 | 2.378e-07 | linear_util.py:343(hashable_partial) |
| 3 | 1.422e-06 | 4.74e-07 | 1.788e-06 | 5.96e-07 | prng.py:85(__hash__) |
| 1 | 1.386e-06 | 1.386e-06 | 2.392e-05 | 2.392e-05 | <__array_function__ internals>:177(tile) |
| 2 | 1.362e-06 | 6.81e-07 | 1.362e-06 | 6.81e-07 | api_util.py:268(<listcomp>) |
| 1 | 1.337e-06 | 1.337e-06 | 0.0001275 | 0.0001275 | lax.py:878(select) |
| 3 | 1.296e-06 | 4.32e-07 | 6.774e-06 | 2.258e-06 | dtypes.py:187(coerce_to_array) |
| 1 | 1.226e-06 | 1.226e-06 | 1.652e-05 | 1.652e-05 | <__array_function__ internals>:177(cumsum) |
| 4 | 1.205e-06 | 3.013e-07 | 1.205e-06 | 3.013e-07 | ~:0(<built-in method numpy.geterrobj>) |
| 1 | 1.198e-06 | 1.198e-06 | 0.0001066 | 0.0001066 | lax_numpy.py:855(_squeeze) |
| 1 | 1.178e-06 | 1.178e-06 | 1.363e-05 | 1.363e-05 | fromnumeric.py:2523(cumsum) |
| 1 | 1.159e-06 | 1.159e-06 | 0.0001085 | 0.0001085 | lax.py:540(bitcast_convert_type) |
| 2 | 1.111e-06 | 5.555e-07 | 2.532e-06 | 1.266e-06 | lax.py:1214(<listcomp>) |
| 4 | 1.061e-06 | 2.653e-07 | 1.061e-06 | 2.653e-07 | lax.py:3147(_reshape_dtype_rule) |
| 3 | 1.027e-06 | 3.423e-07 | 1.027e-06 | 3.423e-07 | shape_base.py:1273(<genexpr>) |
| 2 | 1.022e-06 | 5.11e-07 | 1.386e-06 | 6.93e-07 | ad.py:131(<listcomp>) |
| 2 | 1.009e-06 | 5.045e-07 | 1.009e-06 | 5.045e-07 | ad.py:715(<listcomp>) |
| 8 | 1.001e-06 | 1.251e-07 | 1.001e-06 | 1.251e-07 | partial_eval.py:1486(_contents) |
| 1 | 9.84e-07 | 9.84e-07 | 1.46e-06 | 1.46e-06 | random.py:231(_check_shape) |
| 1 | 9.78e-07 | 9.78e-07 | 5.759e-05 | 5.759e-05 | util.py:280(<listcomp>) |
| 1 | 9.4e-07 | 9.4e-07 | 0.007218 | 0.007218 | prng.py:615(<lambda>) |
| 6 | 9.37e-07 | 1.562e-07 | 9.37e-07 | 1.562e-07 | ~:0(<method 'pop' of 'set' objects>) |
| 2 | 8.97e-07 | 4.485e-07 | 8.97e-07 | 4.485e-07 | core.py:2089(positional) |
| 1 | 8.71e-07 | 8.71e-07 | 0.0001273 | 0.0001273 | lax.py:356(pow) |
| 1 | 8.68e-07 | 8.68e-07 | 8.68e-07 | 8.68e-07 | ~:0(<method 'disable' of '_lsprof.Profiler' objects>) |
| 1 | 8.63e-07 | 8.63e-07 | 2.741e-06 | 2.741e-06 | core.py:2072(__init__) |
| 1 | 8.57e-07 | 8.57e-07 | 2.442e-06 | 2.442e-06 | core.py:1405(has_opaque_dtype) |
| 1 | 8.46e-07 | 8.46e-07 | 0.0001303 | 0.0001303 | lax.py:480(lt) |
| 2 | 8.37e-07 | 4.185e-07 | 8.37e-07 | 4.185e-07 | hlo_helpers.py:65(<listcomp>) |
| 1 | 8.15e-07 | 8.15e-07 | 0.0001623 | 0.0001623 | lax.py:384(bitwise_or) |
| 2 | 8.09e-07 | 4.045e-07 | 8.09e-07 | 4.045e-07 | prng.py:88(__str__) |
| 1 | 7.53e-07 | 7.53e-07 | 0.0001678 | 0.0001678 | lax.py:456(shift_right_logical) |
| 2 | 7.53e-07 | 3.765e-07 | 7.53e-07 | 3.765e-07 | ad.py:716(<listcomp>) |
| 1 | 7.42e-07 | 7.42e-07 | 8.616e-05 | 8.616e-05 | lax.py:1684(<lambda>) |
| 2 | 7.42e-07 | 3.71e-07 | 1.029e-06 | 5.145e-07 | core.py:2081(positional_rank) |
| 1 | 7.15e-07 | 7.15e-07 | 2.15e-06 | 2.15e-06 | core.py:1980(is_empty_shape) |
| 4 | 6.96e-07 | 1.74e-07 | 6.96e-07 | 1.74e-07 | contextlib.py:754(__enter__) |
| 3 | 6.95e-07 | 2.317e-07 | 6.95e-07 | 2.317e-07 | utils.py:42(cast_tree) |
| 2 | 6.81e-07 | 3.405e-07 | 6.81e-07 | 3.405e-07 | lax.py:3252(<genexpr>) |
| 2 | 6.38e-07 | 3.19e-07 | 6.38e-07 | 3.19e-07 | dtypes.py:587(is_python_scalar) |
| 2 | 6.21e-07 | 3.105e-07 | 6.21e-07 | 3.105e-07 | ad.py:110(<genexpr>) |
| 1 | 5.83e-07 | 5.83e-07 | 0.0001139 | 0.0001139 | lax.py:2100(<lambda>) |
| 4 | 5.75e-07 | 1.438e-07 | 5.75e-07 | 1.438e-07 | contextlib.py:757(__exit__) |
| 3 | 5.44e-07 | 1.813e-07 | 5.44e-07 | 1.813e-07 | lax.py:3262(<genexpr>) |
| 1 | 5.11e-07 | 5.11e-07 | 5.11e-07 | 5.11e-07 | prng.py:466(__eq__) |
| 2 | 5.05e-07 | 2.525e-07 | 5.05e-07 | 2.525e-07 | shape_base.py:1263(<genexpr>) |
| 4 | 4.95e-07 | 1.238e-07 | 4.95e-07 | 1.238e-07 | lax.py:1490(<lambda>) |
| 2 | 4.92e-07 | 2.46e-07 | 4.92e-07 | 2.46e-07 | multiarray.py:502(can_cast) |
| 2 | 4.82e-07 | 2.41e-07 | 4.82e-07 | 2.41e-07 | lax.py:3272(<genexpr>) |
| 1 | 4.73e-07 | 4.73e-07 | 6.02e-07 | 6.02e-07 | core.py:2101(named_items) |
| 2 | 4.6e-07 | 2.3e-07 | 4.6e-07 | 2.3e-07 | prng.py:749(random_wrap_lowering) |
| 1 | 4.34e-07 | 4.34e-07 | 4.34e-07 | 4.34e-07 | lax.py:2937(<listcomp>) |
| 3 | 4.03e-07 | 1.343e-07 | 4.03e-07 | 1.343e-07 | ad_util.py:94(<lambda>) |
| 2 | 4.02e-07 | 2.01e-07 | 4.02e-07 | 2.01e-07 | lax.py:4302(<lambda>) |
| 1 | 3.84e-07 | 3.84e-07 | 3.84e-07 | 3.84e-07 | ~:0(<built-in method from_bytes>) |
| 2 | 3.75e-07 | 1.875e-07 | 3.75e-07 | 1.875e-07 | pxla.py:3164(<lambda>) |
| 2 | 2.84e-07 | 1.42e-07 | 2.84e-07 | 1.42e-07 | lax.py:4301(<lambda>) |
| 1 | 2.26e-07 | 2.26e-07 | 2.26e-07 | 2.26e-07 | prng.py:779(random_unwrap_lowering) |
| 1 | 2.12e-07 | 2.12e-07 | 2.12e-07 | 2.12e-07 | tree_util.py:239(<lambda>) |
| 1 | 2.1e-07 | 2.1e-07 | 2.1e-07 | 2.1e-07 | core.py:1981(<genexpr>) |
| 2 | 2.09e-07 | 1.045e-07 | 2.09e-07 | 1.045e-07 | config.py:600(<lambda>) |
| 1 | 1.83e-07 | 1.83e-07 | 1.83e-07 | 1.83e-07 | shape_base.py:1187(_tile_dispatcher) |
| 1 | 1.69e-07 | 1.69e-07 | 1.69e-07 | 1.69e-07 | fromnumeric.py:2519(_cumsum_dispatcher) |
| 1 | 1.45e-07 | 1.45e-07 | 1.45e-07 | 1.45e-07 | mlir.py:1423(<lambda>) |
| 1 | 1.33e-07 | 1.33e-07 | 1.33e-07 | 1.33e-07 | contextlib.py:65(_recreate_cm) |